92 lines
1.7 KiB
Go
92 lines
1.7 KiB
Go
|
package main
|
||
|
|
|
||
|
|
import (
|
||
|
|
"database/sql"
|
||
|
|
"log"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Database to store a single connection to the database and its statistics
|
||
|
|
type Database struct {
|
||
|
|
db *sql.DB
|
||
|
|
driver string
|
||
|
|
dsn string
|
||
|
|
query string
|
||
|
|
reconnect bool
|
||
|
|
queries float64
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewDatabase creates a single connection to the database then returns the struct
|
||
|
|
func NewDatabase(driver string, dsn string, query string, reconnect bool) (*Database, error) {
|
||
|
|
database := &Database{
|
||
|
|
driver: driver,
|
||
|
|
dsn: dsn,
|
||
|
|
query: query,
|
||
|
|
reconnect: reconnect,
|
||
|
|
}
|
||
|
|
err := database.connect()
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return database, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (d *Database) connect() error {
|
||
|
|
db, err := sql.Open(d.driver, d.dsn)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if err = db.Ping(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
d.db = db
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (d *Database) disconnect() error {
|
||
|
|
if d.db != nil {
|
||
|
|
return d.db.Close()
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Run to perform benchmark on a single connection for a duration
|
||
|
|
func (d *Database) Run(duration time.Duration, wg *sync.WaitGroup) error {
|
||
|
|
defer wg.Done()
|
||
|
|
defer d.disconnect()
|
||
|
|
|
||
|
|
// Run single benchmark
|
||
|
|
start := time.Now()
|
||
|
|
end := start.Add(duration)
|
||
|
|
for {
|
||
|
|
if end.Before(time.Now()) {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
result, err := d.db.Query(d.query)
|
||
|
|
if err != nil {
|
||
|
|
log.Fatalf("Failed query the database: %v", err)
|
||
|
|
}
|
||
|
|
if err = result.Close(); err != nil {
|
||
|
|
log.Fatalf("Failed to close query: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
d.queries++
|
||
|
|
if d.reconnect {
|
||
|
|
if err = d.disconnect(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if err = d.connect(); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Queries returns the number of performed queries in the benchmark
|
||
|
|
func (d *Database) Queries() float64 {
|
||
|
|
return d.queries
|
||
|
|
}
|