Files
cyclingbot/db/db.go
2026-04-17 19:38:35 -05:00

522 lines
16 KiB
Go

package db
import (
"context"
"database/sql"
"fmt"
"time"
_ "github.com/lib/pq"
)
type DB struct {
conn *sql.DB
}
type UserStats struct {
UserID string
Username string
TotalKM float64
LogCount int
LastUpdated string
}
type RideLog struct {
ID int64
KM float64
MessageID string
LoggedAt time.Time
}
type ChallengeArchive struct {
Name string
TotalKM float64
Riders int
StartDate time.Time
EndDate time.Time
}
func Open(connStr string) (*DB, error) {
conn, err := sql.Open("postgres", connStr)
if err != nil {
return nil, fmt.Errorf("open db: %w", err)
}
if err := conn.Ping(); err != nil {
return nil, fmt.Errorf("ping db: %w", err)
}
d := &DB{conn: conn}
if err := d.migrate(); err != nil {
return nil, fmt.Errorf("migrate: %w", err)
}
return d, nil
}
func (d *DB) Close() error {
return d.conn.Close()
}
func (d *DB) migrate() error {
_, err := d.conn.Exec(`
CREATE TABLE IF NOT EXISTS distance_logs (
id SERIAL PRIMARY KEY,
guild_id TEXT NOT NULL DEFAULT '',
user_id TEXT NOT NULL,
username TEXT NOT NULL,
km DOUBLE PRECISION NOT NULL,
message_id TEXT NOT NULL UNIQUE,
channel_id TEXT NOT NULL,
logged_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS settings (
guild_id TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
PRIMARY KEY (guild_id, key)
);
CREATE TABLE IF NOT EXISTS challenge_archive (
id SERIAL PRIMARY KEY,
guild_id TEXT NOT NULL,
name TEXT NOT NULL DEFAULT '',
total_km DOUBLE PRECISION NOT NULL,
riders INTEGER NOT NULL,
start_date TIMESTAMPTZ NOT NULL,
end_date TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS user_preferences (
user_id TEXT NOT NULL,
guild_id TEXT NOT NULL,
unit TEXT NOT NULL DEFAULT 'km',
PRIMARY KEY (user_id, guild_id)
);
CREATE TABLE IF NOT EXISTS kudos (
id SERIAL PRIMARY KEY,
guild_id TEXT NOT NULL,
from_user_id TEXT NOT NULL,
from_username TEXT NOT NULL,
to_user_id TEXT NOT NULL,
to_username TEXT NOT NULL,
given_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
);
`)
if err != nil {
return err
}
_, _ = d.conn.Exec(`ALTER TABLE distance_logs ADD COLUMN IF NOT EXISTS guild_id TEXT NOT NULL DEFAULT ''`)
return nil
}
// ── Settings ─────────────────────────────────────────────────────────────────
func (d *DB) GetSetting(ctx context.Context, guildID, key string) (string, bool, error) {
var value string
err := d.conn.QueryRowContext(ctx,
`SELECT value FROM settings WHERE guild_id = $1 AND key = $2`, guildID, key,
).Scan(&value)
if err == sql.ErrNoRows {
return "", false, nil
}
return value, err == nil, err
}
func (d *DB) SetSetting(ctx context.Context, guildID, key, value string) error {
_, err := d.conn.ExecContext(ctx, `
INSERT INTO settings (guild_id, key, value) VALUES ($1, $2, $3)
ON CONFLICT(guild_id, key) DO UPDATE SET value = EXCLUDED.value
`, guildID, key, value)
return err
}
func (d *DB) DeleteSetting(ctx context.Context, guildID, key string) error {
_, err := d.conn.ExecContext(ctx,
`DELETE FROM settings WHERE guild_id = $1 AND key = $2`, guildID, key)
return err
}
// ── Core Logging ──────────────────────────────────────────────────────────────
func (d *DB) AddLog(ctx context.Context, guildID, userID, username, messageID, channelID string, km float64) (bool, error) {
res, err := d.conn.ExecContext(ctx, `
INSERT INTO distance_logs (guild_id, user_id, username, km, message_id, channel_id)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT(message_id) DO NOTHING
`, guildID, userID, username, km, messageID, channelID)
if err != nil {
return false, fmt.Errorf("insert log: %w", err)
}
rows, _ := res.RowsAffected()
return rows > 0, nil
}
func (d *DB) RemoveLog(ctx context.Context, messageID string) (float64, error) {
var km float64
err := d.conn.QueryRowContext(ctx,
`SELECT km FROM distance_logs WHERE message_id = $1`, messageID).Scan(&km)
if err == sql.ErrNoRows {
return 0, nil
}
if err != nil {
return 0, err
}
_, err = d.conn.ExecContext(ctx, `DELETE FROM distance_logs WHERE message_id = $1`, messageID)
return km, err
}
func (d *DB) AdjustKM(ctx context.Context, guildID, userID, username string, km float64) error {
_, err := d.conn.ExecContext(ctx, `
INSERT INTO distance_logs (guild_id, user_id, username, km, message_id, channel_id)
VALUES ($1, $2, $3, $4, 'manual-' || gen_random_uuid()::text, 'admin')
`, guildID, userID, username, km)
return err
}
// ── Stats Queries ─────────────────────────────────────────────────────────────
func (d *DB) GetLeaderboard(ctx context.Context, guildID string, since time.Time, limit int) ([]*UserStats, error) {
q := `SELECT user_id, username, SUM(km), COUNT(*), MAX(logged_at)
FROM distance_logs WHERE guild_id = $1`
args := []interface{}{guildID}
if !since.IsZero() {
args = append(args, since)
q += fmt.Sprintf(` AND logged_at >= $%d`, len(args))
}
args = append(args, limit)
q += fmt.Sprintf(` GROUP BY user_id, username ORDER BY SUM(km) DESC LIMIT $%d`, len(args))
rows, err := d.conn.QueryContext(ctx, q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
var results []*UserStats
for rows.Next() {
s := &UserStats{}
var lastUpdated time.Time
if err := rows.Scan(&s.UserID, &s.Username, &s.TotalKM, &s.LogCount, &lastUpdated); err != nil {
return nil, err
}
s.LastUpdated = lastUpdated.Format(time.RFC3339)
results = append(results, s)
}
return results, rows.Err()
}
func (d *DB) GetTotalKM(ctx context.Context, guildID string, since time.Time) (float64, error) {
q := `SELECT COALESCE(SUM(km), 0) FROM distance_logs WHERE guild_id = $1`
args := []interface{}{guildID}
if !since.IsZero() {
args = append(args, since)
q += fmt.Sprintf(` AND logged_at >= $%d`, len(args))
}
var total float64
err := d.conn.QueryRowContext(ctx, q, args...).Scan(&total)
return total, err
}
func (d *DB) GetUserStats(ctx context.Context, guildID, userID string, since time.Time) (*UserStats, error) {
q := `SELECT user_id, username, COALESCE(SUM(km), 0), COUNT(*), COALESCE(MAX(logged_at)::text, '')
FROM distance_logs WHERE guild_id = $1 AND user_id = $2`
args := []interface{}{guildID, userID}
if !since.IsZero() {
args = append(args, since)
q += fmt.Sprintf(` AND logged_at >= $%d`, len(args))
}
q += ` GROUP BY user_id, username`
s := &UserStats{}
err := d.conn.QueryRowContext(ctx, q, args...).Scan(
&s.UserID, &s.Username, &s.TotalKM, &s.LogCount, &s.LastUpdated)
if err == sql.ErrNoRows {
return &UserStats{UserID: userID}, nil
}
return s, err
}
func (d *DB) GetStatsInRange(ctx context.Context, guildID string, from, to time.Time, limit int) ([]*UserStats, error) {
rows, err := d.conn.QueryContext(ctx, `
SELECT user_id, username, SUM(km), COUNT(*), MAX(logged_at)
FROM distance_logs
WHERE guild_id = $1 AND logged_at >= $2 AND logged_at <= $3
GROUP BY user_id, username
ORDER BY SUM(km) DESC
LIMIT $4
`, guildID, from, to, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var results []*UserStats
for rows.Next() {
s := &UserStats{}
var lastUpdated time.Time
if err := rows.Scan(&s.UserID, &s.Username, &s.TotalKM, &s.LogCount, &lastUpdated); err != nil {
return nil, err
}
s.LastUpdated = lastUpdated.Format(time.RFC3339)
results = append(results, s)
}
return results, rows.Err()
}
func (d *DB) GetYearTotals(ctx context.Context, guildID string) ([]YearTotal, error) {
rows, err := d.conn.QueryContext(ctx, `
SELECT EXTRACT(YEAR FROM logged_at)::INTEGER as yr, SUM(km)
FROM distance_logs
WHERE guild_id = $1
GROUP BY yr
ORDER BY yr DESC
`, guildID)
if err != nil {
return nil, err
}
defer rows.Close()
var results []YearTotal
for rows.Next() {
var yt YearTotal
if err := rows.Scan(&yt.Year, &yt.TotalKM); err != nil {
return nil, err
}
results = append(results, yt)
}
return results, rows.Err()
}
type YearTotal struct {
Year int
TotalKM float64
}
func (d *DB) GetYearlyLeaderboard(ctx context.Context, guildID string, year, limit int) ([]*UserStats, error) {
rows, err := d.conn.QueryContext(ctx, `
SELECT user_id, username, SUM(km), COUNT(*), MAX(logged_at)
FROM distance_logs
WHERE guild_id = $1 AND EXTRACT(YEAR FROM logged_at) = $2
GROUP BY user_id, username
ORDER BY SUM(km) DESC
LIMIT $3
`, guildID, year, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var results []*UserStats
for rows.Next() {
s := &UserStats{}
var lastUpdated time.Time
if err := rows.Scan(&s.UserID, &s.Username, &s.TotalKM, &s.LogCount, &lastUpdated); err != nil {
return nil, err
}
s.LastUpdated = lastUpdated.Format(time.RFC3339)
results = append(results, s)
}
return results, rows.Err()
}
func (d *DB) GetUserYearlyStats(ctx context.Context, guildID, userID string, year int) (*UserStats, error) {
s := &UserStats{UserID: userID}
var lastUpdated sql.NullTime
err := d.conn.QueryRowContext(ctx, `
SELECT user_id, username, COALESCE(SUM(km), 0), COUNT(*), MAX(logged_at)
FROM distance_logs
WHERE guild_id = $1 AND user_id = $2 AND EXTRACT(YEAR FROM logged_at) = $3
GROUP BY user_id, username
`, guildID, userID, year).Scan(&s.UserID, &s.Username, &s.TotalKM, &s.LogCount, &lastUpdated)
if err == sql.ErrNoRows {
return s, nil
}
if lastUpdated.Valid {
s.LastUpdated = lastUpdated.Time.Format(time.RFC3339)
}
return s, err
}
// ── Personal Stats ────────────────────────────────────────────────────────────
func (d *DB) GetUserHistory(ctx context.Context, guildID, userID string, limit int) ([]*RideLog, error) {
rows, err := d.conn.QueryContext(ctx, `
SELECT id, km, message_id, logged_at
FROM distance_logs
WHERE guild_id = $1 AND user_id = $2
ORDER BY logged_at DESC
LIMIT $3
`, guildID, userID, limit)
if err != nil {
return nil, err
}
defer rows.Close()
var logs []*RideLog
for rows.Next() {
l := &RideLog{}
if err := rows.Scan(&l.ID, &l.KM, &l.MessageID, &l.LoggedAt); err != nil {
return nil, err
}
logs = append(logs, l)
}
return logs, rows.Err()
}
func (d *DB) GetUserPB(ctx context.Context, guildID, userID string) (float64, error) {
var pb float64
err := d.conn.QueryRowContext(ctx,
`SELECT COALESCE(MAX(km), 0) FROM distance_logs WHERE guild_id = $1 AND user_id = $2`,
guildID, userID).Scan(&pb)
return pb, err
}
func (d *DB) GetUserStreak(ctx context.Context, userID string) (int, error) {
rows, err := d.conn.QueryContext(ctx, `
SELECT DISTINCT logged_at::date as ride_date
FROM distance_logs
WHERE user_id = $1
ORDER BY ride_date DESC
`, userID)
if err != nil {
return 0, err
}
defer rows.Close()
var dates []time.Time
for rows.Next() {
var t time.Time
if err := rows.Scan(&t); err != nil {
return 0, err
}
dates = append(dates, t)
}
if len(dates) == 0 {
return 0, nil
}
today := time.Now().UTC().Truncate(24 * time.Hour)
if dates[0].Before(today.Add(-24 * time.Hour)) {
return 0, nil
}
streak := 1
for i := 1; i < len(dates); i++ {
if dates[i-1].Sub(dates[i]) == 24*time.Hour {
streak++
} else {
break
}
}
return streak, nil
}
// ── User Preferences ──────────────────────────────────────────────────────────
func (d *DB) GetUserPreference(ctx context.Context, userID, guildID string) (string, error) {
var unit string
err := d.conn.QueryRowContext(ctx,
`SELECT unit FROM user_preferences WHERE user_id = $1 AND guild_id = $2`,
userID, guildID).Scan(&unit)
if err == sql.ErrNoRows {
return "km", nil
}
return unit, err
}
func (d *DB) SetUserPreference(ctx context.Context, userID, guildID, unit string) error {
_, err := d.conn.ExecContext(ctx, `
INSERT INTO user_preferences (user_id, guild_id, unit) VALUES ($1, $2, $3)
ON CONFLICT(user_id, guild_id) DO UPDATE SET unit = EXCLUDED.unit
`, userID, guildID, unit)
return err
}
// ── Kudos ─────────────────────────────────────────────────────────────────────
func (d *DB) GiveKudos(ctx context.Context, guildID, fromUserID, fromUsername, toUserID, toUsername string) error {
_, err := d.conn.ExecContext(ctx, `
INSERT INTO kudos (guild_id, from_user_id, from_username, to_user_id, to_username)
VALUES ($1, $2, $3, $4, $5)
`, guildID, fromUserID, fromUsername, toUserID, toUsername)
return err
}
func (d *DB) GetKudosReceived(ctx context.Context, guildID, toUserID string) (int, error) {
var count int
err := d.conn.QueryRowContext(ctx,
`SELECT COUNT(*) FROM kudos WHERE guild_id = $1 AND to_user_id = $2`,
guildID, toUserID).Scan(&count)
return count, err
}
// ── Challenge Management ──────────────────────────────────────────────────────
func (d *DB) GetChallengeStart(ctx context.Context, guildID string) (time.Time, bool, error) {
val, ok, err := d.GetSetting(ctx, guildID, "challenge_start")
if !ok || err != nil {
return time.Time{}, false, err
}
t, err := time.Parse(time.RFC3339, val)
return t, err == nil, err
}
func (d *DB) ResetChallenge(ctx context.Context, guildID, name string) error {
challengeStart, hasPrev, _ := d.GetChallengeStart(ctx, guildID)
total, err := d.GetTotalKM(ctx, guildID, challengeStart)
if err != nil {
return err
}
entries, err := d.GetLeaderboard(ctx, guildID, challengeStart, 1000)
if err != nil {
return err
}
if hasPrev && total > 0 {
archiveName := name
if archiveName == "" {
archiveName = "Challenge"
}
_, err = d.conn.ExecContext(ctx, `
INSERT INTO challenge_archive (guild_id, name, total_km, riders, start_date)
VALUES ($1, $2, $3, $4, $5)
`, guildID, archiveName, total, len(entries), challengeStart)
if err != nil {
return err
}
}
return d.SetSetting(ctx, guildID, "challenge_start", time.Now().UTC().Format(time.RFC3339))
}
func (d *DB) GetChallengeArchive(ctx context.Context, guildID string) ([]*ChallengeArchive, error) {
rows, err := d.conn.QueryContext(ctx, `
SELECT name, total_km, riders, start_date, end_date
FROM challenge_archive
WHERE guild_id = $1
ORDER BY end_date DESC
`, guildID)
if err != nil {
return nil, err
}
defer rows.Close()
var results []*ChallengeArchive
for rows.Next() {
a := &ChallengeArchive{}
if err := rows.Scan(&a.Name, &a.TotalKM, &a.Riders, &a.StartDate, &a.EndDate); err != nil {
return nil, err
}
results = append(results, a)
}
return results, rows.Err()
}
// ── Admin ─────────────────────────────────────────────────────────────────────
func (d *DB) GetUserLogs(ctx context.Context, guildID, userID string, limit int) ([]*RideLog, error) {
return d.GetUserHistory(ctx, guildID, userID, limit)
}