Skip to content

Commit

Permalink
Specify a custom dial function per config (go-sql-driver#1527)
Browse files Browse the repository at this point in the history
Specify a custom dial function per config instead of using
RegisterDialContext.
  • Loading branch information
aaronjheng authored and methane committed Nov 13, 2024
1 parent aa316e7 commit cdaf43e
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 13 deletions.
31 changes: 18 additions & 13 deletions connector.go
Original file line number Diff line number Diff line change
Expand Up @@ -87,20 +87,25 @@ 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 {
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)
dctx := ctx
if mc.cfg.Timeout > 0 {
var cancel context.CancelFunc
dctx, cancel = context.WithTimeout(ctx, c.cfg.Timeout)
defer cancel()
}

if c.cfg.DialFunc != nil {
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 {
mc.netConn, err = dial(dctx, mc.cfg.Addr)
} else {
nd := net.Dialer{}
mc.netConn, err = nd.DialContext(dctx, mc.cfg.Net, mc.cfg.Addr)
}
}
if err != nil {
return nil, err
Expand Down
2 changes: 2 additions & 0 deletions dsn.go
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ type Config struct {
ReadTimeout time.Duration // I/O read timeout
WriteTimeout time.Duration // I/O write timeout
Logger Logger // Logger
// DialFunc specifies the dial function for creating connections
DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

// boolean fields

Expand Down

0 comments on commit cdaf43e

Please sign in to comment.