-
Notifications
You must be signed in to change notification settings - Fork 0
/
driver.go
100 lines (85 loc) · 3.08 KB
/
driver.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
package database
import (
"context"
"database/sql/driver"
"github.com/icinga/icinga-go-library/backoff"
"github.com/icinga/icinga-go-library/logging"
"github.com/icinga/icinga-go-library/retry"
"github.com/pkg/errors"
"go.uber.org/zap"
"time"
)
// Driver names as automatically registered in the database/sql package by themselves.
const (
MySQL string = "mysql"
PostgreSQL string = "postgres"
)
// OnInitConnFunc can be used to execute post Connect() arbitrary actions.
// It will be called after successfully initiated a new connection using the connector's Connect method.
type OnInitConnFunc func(context.Context, driver.Conn) error
// RetryConnectorCallbacks specifies callbacks that are executed upon certain events.
type RetryConnectorCallbacks struct {
OnInitConn OnInitConnFunc
OnRetryableError retry.OnRetryableErrorFunc
OnSuccess retry.OnSuccessFunc
}
// RetryConnector wraps driver.Connector with retry logic.
type RetryConnector struct {
driver.Connector
logger *logging.Logger
callbacks RetryConnectorCallbacks
}
// NewConnector creates a fully initialized RetryConnector from the given args.
func NewConnector(c driver.Connector, logger *logging.Logger, callbacks RetryConnectorCallbacks) *RetryConnector {
return &RetryConnector{Connector: c, logger: logger, callbacks: callbacks}
}
// Connect implements part of the driver.Connector interface.
func (c RetryConnector) Connect(ctx context.Context) (driver.Conn, error) {
var conn driver.Conn
err := errors.Wrap(retry.WithBackoff(
ctx,
func(ctx context.Context) (err error) {
conn, err = c.Connector.Connect(ctx)
if err == nil && c.callbacks.OnInitConn != nil {
if err = c.callbacks.OnInitConn(ctx, conn); err != nil {
// We're going to retry this, so just don't bother whether Close() fails!
_ = conn.Close()
}
}
return
},
retry.Retryable,
backoff.NewExponentialWithJitter(128*time.Millisecond, 1*time.Minute),
retry.Settings{
Timeout: retry.DefaultTimeout,
OnRetryableError: func(elapsed time.Duration, attempt uint64, err, lastErr error) {
if c.callbacks.OnRetryableError != nil {
c.callbacks.OnRetryableError(elapsed, attempt, err, lastErr)
}
if lastErr == nil || err.Error() != lastErr.Error() {
c.logger.Warnw("Can't connect to database. Retrying", zap.Error(err))
}
},
OnSuccess: func(elapsed time.Duration, attempt uint64, lastErr error) {
if c.callbacks.OnSuccess != nil {
c.callbacks.OnSuccess(elapsed, attempt, lastErr)
}
if attempt > 1 {
c.logger.Infow("Reconnected to database",
zap.Duration("after", elapsed), zap.Uint64("attempts", attempt))
}
},
},
), "can't connect to database")
return conn, err
}
// Driver implements part of the driver.Connector interface.
func (c RetryConnector) Driver() driver.Driver {
return c.Connector.Driver()
}
// MysqlFuncLogger is an adapter that allows ordinary functions to be used as a logger for mysql.SetLogger.
type MysqlFuncLogger func(v ...interface{})
// Print implements the mysql.Logger interface.
func (log MysqlFuncLogger) Print(v ...interface{}) {
log(v)
}