Specify a custom dial function per config

Author: aaronjheng

connector.go

@@ -80,7 +80,16 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	dialsLock.RLock()
 	dial, ok := dials[mc.cfg.Net]
 	dialsLock.RUnlock()
-	if ok {
+
+	if c.cfg.DialFunc != nil {
+		dctx := ctx
+		if mc.cfg.Timeout > 0 {
+			var cancel context.CancelFunc
+			dctx, cancel = context.WithTimeout(ctx, c.cfg.Timeout)
+			defer cancel()
+		}
+		mc.netConn, err = c.cfg.DialFunc(dctx, mc.cfg.Addr)
+	} else if ok {
 		dctx := ctx
 		if mc.cfg.Timeout > 0 {
 			var cancel context.CancelFunc

dsn.go

@@ -52,6 +52,7 @@ type Config struct {
 	ReadTimeout          time.Duration     // I/O read timeout
 	WriteTimeout         time.Duration     // I/O write timeout
 	Logger               Logger            // Logger
+	DialFunc             DialContextFunc   // Specify a custom dial function per config instead of using RegisterDialContext
 
 	AllowAllFiles            bool // Allow all files to be used with LOAD DATA LOCAL INFILE
 	AllowCleartextPasswords  bool // Allows the cleartext client side plugin