64 lines
1.4 KiB
Go
64 lines
1.4 KiB
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()
|
||
|
|
}
|
||
|
|
|
||
|
|
type hostTotals struct {
|
||
|
|
Upload int64
|
||
|
|
Download int64
|
||
|
|
}
|