Driver context

pkg/driver/driver.go

@@ -18,31 +18,58 @@ import (
 
 var timeout = time.Minute * 5
 
-// TODO(el): Support DriverContext.
-type Driver struct {
-	Driver driver.Driver
-	Logger *zap.SugaredLogger
+// RetryConnector wraps driver.Connector with retry logic.
+type RetryConnector struct {
+	driver.Connector
+	driver Driver
 }
 
-// TODO(el): Test DNS.
-func (d Driver) Open(dsn string) (c driver.Conn, err error) {
+// Connect implements part of the driver.Connector interface.
+func (c RetryConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	var conn driver.Conn
 	var logFirstError sync.Once
-	err = errors.Wrap(retry.WithBackoff(
-		context.Background(),
-		func(context.Context) (err error) {
-			c, err = d.Driver.Open(dsn)
+	err := errors.Wrap(retry.WithBackoff(
+		ctx,
+		func(ctx context.Context) (err error) {
+			conn, err = c.Connector.Connect(ctx)
+
 			logFirstError.Do(func() {
 				if err != nil {
-					d.Logger.Warnw("Can't connect to database. Retrying", zap.Error(err))
+					c.driver.Logger.Warnw("Can't connect to database. Retrying", zap.Error(err))
 				}
 			})
+
 			return
 		},
 		shouldRetry,
 		backoff.NewExponentialWithJitter(time.Millisecond*128, time.Minute*1),
 		timeout,
 	), "can't connect to database")
-	return
+	return conn, err
+}
+
+// Driver implements part of the driver.Connector interface.
+func (c RetryConnector) Driver() driver.Driver {
+	return c.driver
+}
+
+// Driver wraps a driver.Driver that also must implement driver.DriverContext with logging capabilities and provides our RetryConnector.
+type Driver struct {
+	ctxDriver
+	Logger *zap.SugaredLogger
+}
+
+// OpenConnector implements the DriverContext interface.
+func (d Driver) OpenConnector(name string) (driver.Connector, error) {
+	c, err := d.ctxDriver.OpenConnector(name)
+	if err != nil {
+		return nil, err
+	}
+
+	return &RetryConnector{
+		driver:    d,
+		Connector: c,
+	}, nil
 }
 
 func shouldRetry(err error) bool {
@@ -68,7 +95,13 @@ func shouldRetry(err error) bool {
 }
 
 func Register(logger *zap.SugaredLogger) {
-	sql.Register("icingadb-mysql", &Driver{Driver: &mysql.MySQLDriver{}, Logger: logger})
+	sql.Register("icingadb-mysql", &Driver{ctxDriver: &mysql.MySQLDriver{}, Logger: logger})
 	// TODO(el): Don't discard but hide?
 	_ = mysql.SetLogger(log.New(ioutil.Discard, "", 0))
 }
+
+// ctxDriver helps ensure that we only support drivers that implement driver.Driver and driver.DriverContext.
+type ctxDriver interface {
+	driver.Driver
+	driver.DriverContext
+}