Havline/internal/traffic/store.go

package traffic

import (
	"context"
	"database/sql"
	"strings"
	"time"
)

type Store struct {
	db *sql.DB
}

func NewStore(db *sql.DB) *Store {
	return &Store{db: db}
}

func (s *Store) LoadTotals(ctx context.Context) (map[string]hostTotals, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT host, upload_bytes, download_bytes FROM proxy_traffic`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := map[string]hostTotals{}
	for rows.Next() {
		var host string
		var up, down int64
		if err := rows.Scan(&host, &up, &down); err != nil {
			return nil, err
		}
		key := strings.ToLower(strings.TrimSpace(host))
		out[key] = hostTotals{Upload: up, Download: down}
	}
	return out, rows.Err()
}

func (s *Store) SaveTotals(ctx context.Context, hosts map[string]hostTotals) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	for host, totals := range hosts {
		_, err := tx.ExecContext(ctx, `
			INSERT INTO proxy_traffic(host, upload_bytes, download_bytes, updated_at)
			VALUES (?, ?, ?, ?)
			ON CONFLICT(host) DO UPDATE SET
				upload_bytes = excluded.upload_bytes,
				download_bytes = excluded.download_bytes,
				updated_at = excluded.updated_at
		`, host, totals.Upload, totals.Download, time.Now().UTC().Format(time.RFC3339))
		if err != nil {
			_ = tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

// DeleteHosts 删除指定 host:port 的累计流量。
// 规则被删掉后这些行就成了孤儿：既没有界面能查到，又会在同名规则重建时被当成本规则的既有流量。
func (s *Store) DeleteHosts(ctx context.Context, keys []string) error {
	if len(keys) == 0 {
		return nil
	}
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	for _, key := range keys {
		if _, err := tx.ExecContext(ctx,
			`DELETE FROM proxy_traffic WHERE host = ?`, strings.ToLower(strings.TrimSpace(key))); err != nil {
			_ = tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

type hostTotals struct {
	Upload   int64
	Download int64
}