71 lines
1.8 KiB
Go
71 lines
1.8 KiB
Go
|
package firefly
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"flag"
|
||
|
|
"github.com/jackc/pgx/v4"
|
||
|
|
"github.com/jackc/pgx/v4/pgxpool"
|
||
|
|
"github.com/sirupsen/logrus"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
pgsqlConn = flag.String("pgsql", "", "Connection spec for postgresql database")
|
||
|
|
pool *pgxpool.Pool
|
||
|
|
)
|
||
|
|
|
||
|
|
func Initialize() {
|
||
|
|
var err error
|
||
|
|
config, err := pgxpool.ParseConfig(*pgsqlConn)
|
||
|
|
if err != nil {
|
||
|
|
panic(err)
|
||
|
|
}
|
||
|
|
if config.MaxConns == 0 {
|
||
|
|
config.MaxConns = 8
|
||
|
|
}
|
||
|
|
if pool, err = pgxpool.Connect(context.Background(), *pgsqlConn); err != nil {
|
||
|
|
panic(err)
|
||
|
|
}
|
||
|
|
if conn, err := pool.Acquire(context.Background()); err != nil {
|
||
|
|
panic(err)
|
||
|
|
} else {
|
||
|
|
conn.Release()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func AddSpeedTestResult(ip string, speed map[string]float64) {
|
||
|
|
conn, err := pool.Acquire(context.Background())
|
||
|
|
if err != nil {
|
||
|
|
logrus.WithError(err).Error("Unable to get database connection")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
defer conn.Release()
|
||
|
|
tx, err := conn.Begin(context.Background())
|
||
|
|
if err != nil {
|
||
|
|
logrus.WithError(err).Error("Unable to start database transaction")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
row := tx.QueryRow(context.Background(),
|
||
|
|
"INSERT INTO speedtest_entry (id, ip, timestamp) VALUES (DEFAULT, $1, now()) RETURNING id", ip)
|
||
|
|
var id int64
|
||
|
|
if err = row.Scan(&id); err != nil {
|
||
|
|
logrus.WithError(err).Error("Unable to insert new speedtest result entry")
|
||
|
|
_ = tx.Rollback(context.Background())
|
||
|
|
return
|
||
|
|
}
|
||
|
|
batch := &pgx.Batch{}
|
||
|
|
for name, value := range speed {
|
||
|
|
batch.Queue("INSERT INTO speedtest_result (entry_id, node, speed) VALUES ($1, $2, $3)",
|
||
|
|
id, name, value)
|
||
|
|
}
|
||
|
|
results := tx.SendBatch(context.Background(), batch)
|
||
|
|
if err = results.Close(); err != nil {
|
||
|
|
logrus.WithError(err).Error("Unable to insert speedtest results")
|
||
|
|
_ = tx.Rollback(context.Background())
|
||
|
|
return
|
||
|
|
}
|
||
|
|
if err = tx.Commit(context.Background()); err != nil {
|
||
|
|
logrus.WithError(err).Error("Unable to commit the database transaction")
|
||
|
|
_ = tx.Rollback(context.Background())
|
||
|
|
}
|
||
|
|
}
|