package k8s

import (
	"context"
	"net"
	"sync"

	"github.com/pingcap/test-infra/sdk/pkg/proxy"

	"github.com/go-sql-driver/mysql"
	netproxy "golang.org/x/net/proxy"
)

var (
	setMySQLProxyOnce sync.Once
)

func init() {
	dialer := proxy.MySQLProxyDialer()

	setMySQLProxyOnce.Do(func() {
		mysql.RegisterDialContext("tcp", func(ctx context.Context, addr string) (net.Conn, error) {
			if xd, ok := dialer.(netproxy.ContextDialer); ok {
				return xd.DialContext(ctx, "tcp", addr)
			}
			return dialer.Dial("tcp", addr)
		})
	})
}
