package main
import (
"database/sql"
"fmt"
"strings"
"time"
_ "github.com/go-sql-driver/mysql"
)
type DB struct {
conn *sql.DB
}
type Server struct {
ID int
URL string
Name string
AnonName string
Active bool
Insecure bool
ExpectedCode int
ExpectedBody string
CreatedAt time.Time
}
type Agent struct {
ID int
Name string
AnonName string
Token string
LastIP string
LastSeen *time.Time
CreatedAt time.Time
}
type CheckResult struct {
AgentID int
ServerID int
CheckedAt time.Time
Success bool
HTTPCode int
ConnectMS int
TTFBMS int
TotalMS int
}
type ServerStat struct {
Total int
Successes int
AvgTTFBMS int
}
type ChartPoint struct {
T time.Time
Success bool
TotalMS int
TTFBMS int
AgentID int
}
func ensureParseTime(dsn string) string {
if strings.Contains(dsn, "parseTime=") {
return dsn
}
if strings.Contains(dsn, "?") {
return dsn + "&parseTime=true"
}
return dsn + "?parseTime=true"
}
func NewDB(dsn string) (*DB, error) {
conn, err := sql.Open("mysql", ensureParseTime(dsn))
if err != nil {
return nil, err
}
conn.SetMaxOpenConns(10)
conn.SetMaxIdleConns(5)
if err := conn.Ping(); err != nil {
return nil, fmt.Errorf("ping: %w", err)
}
d := &DB{conn: conn}
if err := d.initSchema(); err != nil {
return nil, fmt.Errorf("initSchema: %w", err)
}
return d, nil
}
func (d *DB) initSchema() error {
stmts := []string{
`CREATE TABLE IF NOT EXISTS servers (
id INT AUTO_INCREMENT PRIMARY KEY,
url VARCHAR(512) NOT NULL UNIQUE,
name VARCHAR(255) NOT NULL DEFAULT '',
anon_name VARCHAR(255) NOT NULL DEFAULT '',
active TINYINT(1) NOT NULL DEFAULT 1,
insecure TINYINT(1) NOT NULL DEFAULT 0,
expected_code INT NOT NULL DEFAULT 200,
expected_body VARCHAR(512) NOT NULL DEFAULT '',
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
`CREATE TABLE IF NOT EXISTS agents (
id INT AUTO_INCREMENT PRIMARY KEY,
name VARCHAR(255) NOT NULL,
anon_name VARCHAR(255) NOT NULL DEFAULT '',
token VARCHAR(128) NOT NULL UNIQUE,
last_ip VARCHAR(45),
last_seen DATETIME NULL,
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
}
for _, s := range stmts {
if _, err := d.conn.Exec(s); err != nil {
return err
}
}
if err := d.migrateAddColumn("servers", "anon_name", "VARCHAR(255) NOT NULL DEFAULT '' AFTER name"); err != nil {
return err
}
if err := d.migrateAddColumn("agents", "anon_name", "VARCHAR(255) NOT NULL DEFAULT '' AFTER name"); err != nil {
return err
}
if err := d.migrateAddColumn("servers", "insecure", "TINYINT(1) NOT NULL DEFAULT 0 AFTER active"); err != nil {
return err
}
if err := d.migrateAddColumn("servers", "expected_code", "INT NOT NULL DEFAULT 200 AFTER insecure"); err != nil {
return err
}
if err := d.migrateAddColumn("servers", "expected_body", "VARCHAR(512) NOT NULL DEFAULT '' AFTER expected_code"); err != nil {
return err
}
return d.EnsureResultsTable(time.Now())
}
// migrateAddColumn adds a column to a table only if it doesn't already exist,
// checking information_schema to avoid any locking on already-migrated tables.
func (d *DB) migrateAddColumn(table, column, definition string) error {
var count int
err := d.conn.QueryRow(
`SELECT COUNT(*) FROM information_schema.COLUMNS
WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = ? AND COLUMN_NAME = ?`,
table, column,
).Scan(&count)
if err != nil {
return err
}
if count > 0 {
return nil // already exists
}
_, err = d.conn.Exec(fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", table, column, definition))
return err
}
func ResultsTableName(t time.Time) string {
return fmt.Sprintf("results_%04d_%02d", t.Year(), t.Month())
}
func (d *DB) EnsureResultsTable(t time.Time) error {
name := ResultsTableName(t)
q := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
id BIGINT AUTO_INCREMENT PRIMARY KEY,
agent_id INT NOT NULL,
server_id INT NOT NULL,
checked_at DATETIME NOT NULL,
success TINYINT(1) NOT NULL,
http_code INT NOT NULL DEFAULT 0,
connect_ms INT NOT NULL DEFAULT 0,
ttfb_ms INT NOT NULL DEFAULT 0,
total_ms INT NOT NULL DEFAULT 0,
INDEX idx_server_time (server_id, checked_at),
INDEX idx_agent_time (agent_id, checked_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`, name)
_, err := d.conn.Exec(q)
return err
}
func (d *DB) resultsTables() ([]string, error) {
rows, err := d.conn.Query("SHOW TABLES LIKE 'results_%'")
if err != nil {
return nil, err
}
defer rows.Close()
var tables []string
for rows.Next() {
var t string
if err := rows.Scan(&t); err != nil {
return nil, err
}
tables = append(tables, t)
}
return tables, rows.Err()
}
func (d *DB) ListServers(activeOnly bool) ([]Server, error) {
q := "SELECT id, url, name, anon_name, active, insecure, expected_code, expected_body, created_at FROM servers"
if activeOnly {
q += " WHERE active = 1"
}
q += " ORDER BY id"
rows, err := d.conn.Query(q)
if err != nil {
return nil, err
}
defer rows.Close()
var out []Server
for rows.Next() {
var s Server
if err := rows.Scan(&s.ID, &s.URL, &s.Name, &s.AnonName, &s.Active, &s.Insecure, &s.ExpectedCode, &s.ExpectedBody, &s.CreatedAt); err != nil {
return nil, err
}
out = append(out, s)
}
return out, rows.Err()
}
func (d *DB) AddServer(url, name, anonName string, insecure bool, expectedCode int, expectedBody string) error {
if expectedCode == 0 {
expectedCode = 200
}
_, err := d.conn.Exec(
"INSERT INTO servers (url, name, anon_name, insecure, expected_code, expected_body) VALUES (?, ?, ?, ?, ?, ?)",
url, name, anonName, insecure, expectedCode, expectedBody,
)
return err
}
func (d *DB) SetServerActive(id int, active bool) error {
v := 0
if active {
v = 1
}
_, err := d.conn.Exec("UPDATE servers SET active = ? WHERE id = ?", v, id)
return err
}
func (d *DB) DeleteServer(id int) error {
_, err := d.conn.Exec("DELETE FROM servers WHERE id = ?", id)
return err
}
func (d *DB) ListAgents() ([]Agent, error) {
rows, err := d.conn.Query(
"SELECT id, name, anon_name, token, COALESCE(last_ip,''), last_seen, created_at FROM agents ORDER BY id",
)
if err != nil {
return nil, err
}
defer rows.Close()
var out []Agent
for rows.Next() {
var a Agent
if err := rows.Scan(&a.ID, &a.Name, &a.AnonName, &a.Token, &a.LastIP, &a.LastSeen, &a.CreatedAt); err != nil {
return nil, err
}
out = append(out, a)
}
return out, rows.Err()
}
func (d *DB) AddAgent(name, anonName, token string) error {
_, err := d.conn.Exec("INSERT INTO agents (name, anon_name, token) VALUES (?, ?, ?)", name, anonName, token)
return err
}
func (d *DB) DeleteAgent(id int) error {
_, err := d.conn.Exec("DELETE FROM agents WHERE id = ?", id)
return err
}
func (d *DB) AgentByToken(token string) (*Agent, error) {
var a Agent
err := d.conn.QueryRow(
"SELECT id, name, anon_name, token, COALESCE(last_ip,''), last_seen, created_at FROM agents WHERE token = ?", token,
).Scan(&a.ID, &a.Name, &a.AnonName, &a.Token, &a.LastIP, &a.LastSeen, &a.CreatedAt)
if err == sql.ErrNoRows {
return nil, nil
}
if err != nil {
return nil, err
}
return &a, nil
}
func (d *DB) UpdateAgentSeen(id int, ip string) error {
_, err := d.conn.Exec("UPDATE agents SET last_ip = ?, last_seen = NOW() WHERE id = ?", ip, id)
return err
}
func (d *DB) SaveResults(results []CheckResult) error {
if len(results) == 0 {
return nil
}
now := time.Now()
if err := d.EnsureResultsTable(now); err != nil {
return err
}
table := ResultsTableName(now)
tx, err := d.conn.Begin()
if err != nil {
return err
}
defer tx.Rollback()
stmt, err := tx.Prepare(fmt.Sprintf(
"INSERT INTO %s (agent_id, server_id, checked_at, success, http_code, connect_ms, ttfb_ms, total_ms) VALUES (?,?,?,?,?,?,?,?)",
table,
))
if err != nil {
return err
}
defer stmt.Close()
for _, r := range results {
t := r.CheckedAt
if t.IsZero() {
t = now
}
if _, err := stmt.Exec(r.AgentID, r.ServerID, t, r.Success, r.HTTPCode, r.ConnectMS, r.TTFBMS, r.TotalMS); err != nil {
return err
}
}
return tx.Commit()
}
func (d *DB) ServerStats(since time.Time) (map[int]ServerStat, error) {
tables, err := d.resultsTables()
if err != nil {
return nil, err
}
out := map[int]ServerStat{}
for _, tbl := range tables {
q := fmt.Sprintf(
"SELECT server_id, COUNT(*), SUM(success), AVG(NULLIF(ttfb_ms,0)) FROM %s WHERE checked_at >= ? GROUP BY server_id",
tbl,
)
rows, err := d.conn.Query(q, since)
if err != nil {
continue
}
for rows.Next() {
var sid, total int
var succ sql.NullInt64
var avgTTFB sql.NullFloat64
if err := rows.Scan(&sid, &total, &succ, &avgTTFB); err != nil {
rows.Close()
return nil, err
}
cur := out[sid]
cur.Total += total
cur.Successes += int(succ.Int64)
// weighted average across tables
if avgTTFB.Valid && total > 0 {
cur.AvgTTFBMS = int((float64(cur.AvgTTFBMS*cur.Total) + avgTTFB.Float64*float64(total)) / float64(cur.Total+total))
}
out[sid] = cur
}
rows.Close()
}
return out, nil
}
// ServerChartData returns chart points for a server. agentID=0 means all agents.
func (d *DB) ServerChartData(serverID int, agentID int, since time.Time) ([]ChartPoint, error) {
tables, err := d.resultsTables()
if err != nil {
return nil, err
}
var points []ChartPoint
for _, tbl := range tables {
var q string
var args []interface{}
if agentID > 0 {
q = fmt.Sprintf(
"SELECT checked_at, success, total_ms, ttfb_ms, agent_id FROM %s WHERE server_id = ? AND agent_id = ? AND checked_at >= ? ORDER BY checked_at",
tbl,
)
args = []interface{}{serverID, agentID, since}
} else {
q = fmt.Sprintf(
"SELECT checked_at, success, total_ms, ttfb_ms, agent_id FROM %s WHERE server_id = ? AND checked_at >= ? ORDER BY checked_at",
tbl,
)
args = []interface{}{serverID, since}
}
rows, err := d.conn.Query(q, args...)
if err != nil {
continue
}
for rows.Next() {
var p ChartPoint
if err := rows.Scan(&p.T, &p.Success, &p.TotalMS, &p.TTFBMS, &p.AgentID); err != nil {
rows.Close()
return nil, err
}
points = append(points, p)
}
rows.Close()
}
return points, nil
}
func (d *DB) GetLatestServerResult(serverID int) (*CheckResult, error) {
tables, err := d.resultsTables()
if err != nil {
return nil, err
}
var latest *CheckResult
latestTime := time.Time{}
for _, tbl := range tables {
query := fmt.Sprintf(`
SELECT agent_id, checked_at, success, http_code, connect_ms, ttfb_ms, total_ms
FROM %s
WHERE server_id = ?
ORDER BY checked_at DESC
LIMIT 1`, tbl)
var cr CheckResult
err := d.conn.QueryRow(query, serverID).Scan(
&cr.AgentID, &cr.CheckedAt, &cr.Success, &cr.HTTPCode,
&cr.ConnectMS, &cr.TTFBMS, &cr.TotalMS,
)
if err == nil {
if latest == nil || cr.CheckedAt.After(latestTime) {
latest = &cr
latestTime = cr.CheckedAt
}
}
}
return latest, nil
}
