3b3e95a413
* feat(metrics): grand update * feat(metrics): fixed missing part in ee tracer * feat(assets): added missing arg * feat(metrics): fixed naming problems
109 lines
2.5 KiB
Go
109 lines
2.5 KiB
Go
package postgres
|
|
|
|
import (
|
|
"bytes"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"openreplay/backend/pkg/db/postgres/pool"
|
|
"openreplay/backend/pkg/metrics/database"
|
|
)
|
|
|
|
const (
|
|
insertPrefix = `INSERT INTO `
|
|
insertValues = ` VALUES `
|
|
insertSuffix = ` ON CONFLICT DO NOTHING;`
|
|
)
|
|
|
|
type Bulk interface {
|
|
Append(args ...interface{}) error
|
|
Send() error
|
|
Table() string
|
|
}
|
|
|
|
type bulkImpl struct {
|
|
conn pool.Pool
|
|
metrics database.Database
|
|
table string
|
|
columns string
|
|
template string
|
|
setSize int
|
|
sizeLimit int
|
|
values []interface{}
|
|
}
|
|
|
|
func (b *bulkImpl) Append(args ...interface{}) error {
|
|
if len(args) != b.setSize {
|
|
return fmt.Errorf("wrong number of arguments, waited: %d, got: %d", b.setSize, len(args))
|
|
}
|
|
b.values = append(b.values, args...)
|
|
if len(b.values)/b.setSize >= b.sizeLimit {
|
|
return b.send()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (b *bulkImpl) Send() error {
|
|
if len(b.values) == 0 {
|
|
return nil
|
|
}
|
|
return b.send()
|
|
}
|
|
|
|
func (b *bulkImpl) Table() string {
|
|
return b.table
|
|
}
|
|
|
|
func (b *bulkImpl) send() error {
|
|
start := time.Now()
|
|
size := len(b.values) / b.setSize
|
|
request := bytes.NewBufferString(insertPrefix + b.table + b.columns + insertValues)
|
|
args := make([]interface{}, b.setSize)
|
|
for i := 0; i < len(b.values)/b.setSize; i++ {
|
|
for j := 0; j < b.setSize; j++ {
|
|
args[j] = i*b.setSize + j + 1
|
|
}
|
|
if i > 0 {
|
|
request.WriteByte(',')
|
|
}
|
|
request.WriteString(fmt.Sprintf(b.template, args...))
|
|
}
|
|
request.WriteString(insertSuffix)
|
|
err := b.conn.Exec(request.String(), b.values...)
|
|
b.values = make([]interface{}, 0, b.setSize*b.sizeLimit)
|
|
if err != nil {
|
|
return fmt.Errorf("send bulk err: %s", err)
|
|
}
|
|
// Save bulk metrics
|
|
b.metrics.RecordBulkElements(float64(size), "pg", b.table)
|
|
b.metrics.RecordBulkInsertDuration(float64(time.Now().Sub(start).Milliseconds()), "pg", b.table)
|
|
return nil
|
|
}
|
|
|
|
func NewBulk(conn pool.Pool, metrics database.Database, table, columns, template string, setSize, sizeLimit int) (Bulk, error) {
|
|
switch {
|
|
case conn == nil:
|
|
return nil, errors.New("db conn is empty")
|
|
case table == "":
|
|
return nil, errors.New("table is empty")
|
|
case columns == "":
|
|
return nil, errors.New("columns is empty")
|
|
case template == "":
|
|
return nil, errors.New("template is empty")
|
|
case setSize <= 0:
|
|
return nil, errors.New("set size is wrong")
|
|
case sizeLimit <= 0:
|
|
return nil, errors.New("size limit is wrong")
|
|
}
|
|
return &bulkImpl{
|
|
conn: conn,
|
|
metrics: metrics,
|
|
table: table,
|
|
columns: columns,
|
|
template: template,
|
|
setSize: setSize,
|
|
sizeLimit: sizeLimit,
|
|
values: make([]interface{}, 0, setSize*sizeLimit),
|
|
}, nil
|
|
}
|