Specify a custom dial function per config

go-sql-driver · Oct 14, 2024 · 63f0c4b · 63f0c4b
1 parent 00dc21a
commit 63f0c4b
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 25 deletions.
diff --git a/connector.go b/connector.go
@@ -87,20 +87,30 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	mc.parseTime = mc.cfg.ParseTime
 
 	// Connect to Server
-	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 = dial(dctx, mc.cfg.Addr)
+		mc.netConn, err = c.cfg.DialFunc(dctx, mc.cfg.Net, mc.cfg.Addr)
 	} else {
-		nd := net.Dialer{Timeout: mc.cfg.Timeout}
-		mc.netConn, err = nd.DialContext(ctx, mc.cfg.Net, mc.cfg.Addr)
+		dialsLock.RLock()
+		dial, ok := dials[mc.cfg.Net]
+		dialsLock.RUnlock()
+		if ok {
+			dctx := ctx
+			if mc.cfg.Timeout > 0 {
+				var cancel context.CancelFunc
+				dctx, cancel = context.WithTimeout(ctx, c.cfg.Timeout)
+				defer cancel()
+			}
+			mc.netConn, err = dial(dctx, mc.cfg.Addr)
+		} else {
+			nd := net.Dialer{Timeout: mc.cfg.Timeout}
+			mc.netConn, err = nd.DialContext(ctx, mc.cfg.Net, mc.cfg.Addr)
+		}
 	}
 	if err != nil {
 		return nil, err

diff --git a/dsn.go b/dsn.go
@@ -37,24 +37,25 @@ var (
 type Config struct {
 	// non boolean fields
 
-	User                 string            // Username
-	Passwd               string            // Password (requires User)
-	Net                  string            // Network (e.g. "tcp", "tcp6", "unix". default: "tcp")
-	Addr                 string            // Address (default: "127.0.0.1:3306" for "tcp" and "/tmp/mysql.sock" for "unix")
-	DBName               string            // Database name
-	Params               map[string]string // Connection parameters
-	ConnectionAttributes string            // Connection Attributes, comma-delimited string of user-defined "key:value" pairs
-	charsets             []string          // Connection charset. When set, this will be set in SET NAMES <charset> query
-	Collation            string            // Connection collation. When set, this will be set in SET NAMES <charset> COLLATE <collation> query
-	Loc                  *time.Location    // Location for time.Time values
-	MaxAllowedPacket     int               // Max packet size allowed
-	ServerPubKey         string            // Server public key name
-	TLSConfig            string            // TLS configuration name
-	TLS                  *tls.Config       // TLS configuration, its priority is higher than TLSConfig
-	Timeout              time.Duration     // Dial timeout
-	ReadTimeout          time.Duration     // I/O read timeout
-	WriteTimeout         time.Duration     // I/O write timeout
-	Logger               Logger            // Logger
+	User                 string                                                            // Username
+	Passwd               string                                                            // Password (requires User)
+	Net                  string                                                            // Network (e.g. "tcp", "tcp6", "unix". default: "tcp")
+	Addr                 string                                                            // Address (default: "127.0.0.1:3306" for "tcp" and "/tmp/mysql.sock" for "unix")
+	DBName               string                                                            // Database name
+	Params               map[string]string                                                 // Connection parameters
+	ConnectionAttributes string                                                            // Connection Attributes, comma-delimited string of user-defined "key:value" pairs
+	charsets             []string                                                          // Connection charset. When set, this will be set in SET NAMES <charset> query
+	Collation            string                                                            // Connection collation. When set, this will be set in SET NAMES <charset> COLLATE <collation> query
+	Loc                  *time.Location                                                    // Location for time.Time values
+	MaxAllowedPacket     int                                                               // Max packet size allowed
+	ServerPubKey         string                                                            // Server public key name
+	TLSConfig            string                                                            // TLS configuration name
+	TLS                  *tls.Config                                                       // TLS configuration, its priority is higher than TLSConfig
+	Timeout              time.Duration                                                     // Dial timeout
+	ReadTimeout          time.Duration                                                     // I/O read timeout
+	WriteTimeout         time.Duration                                                     // I/O write timeout
+	Logger               Logger                                                            // Logger
+	DialFunc             func(ctx context.Context, network, addr string) (net.Conn, error) // Specifies the dial function for creating connections
 
 	// boolean fields