ResolvableSelector supports parse addr and extractor options

dsn_selector.go

@@ -17,6 +17,7 @@ package dsn
 import (
 	"errors"
 	"fmt"
+	"net"
 	"sync"
 	"time"
 
@@ -32,12 +33,27 @@ func init() {
 }
 
 // DefaultSelector dsn default selector
-var DefaultSelector = &DsnSelector{dsns: make(map[string]*registry.Node)}
+var DefaultSelector = NewDsnSelector(false)
+
+// NewDsnSelector creates a new dsn selector.
+func NewDsnSelector(keepAddrUnresolved bool) *DsnSelector {
+	var parseAddr func(network, address string) net.Addr
+	if keepAddrUnresolved {
+		parseAddr = func(network, address string) net.Addr {
+			return unresolvedAddr{network: network, address: address}
+		}
+	}
+	return &DsnSelector{
+		parseAddr: parseAddr,
+		dsns:      make(map[string]*registry.Node),
+	}
+}
 
 // DsnSelector returns original service name node, with memory cache
 type DsnSelector struct {
-	dsns map[string]*registry.Node
-	lock sync.RWMutex
+	parseAddr func(network, address string) net.Addr
+	dsns      map[string]*registry.Node
+	lock      sync.RWMutex
 }
 
 // Select selects address from dsn://user:passwd@tcp(ip:port)/db
@@ -60,6 +76,7 @@ func (s *DsnSelector) Select(serviceName string, opt ...selector.Option) (*regis
 	node = &registry.Node{
 		ServiceName: serviceName,
 		Address:     serviceName,
+		ParseAddr:   s.parseAddr,
 	}
 	s.dsns[serviceName] = node
 	return node, nil
@@ -70,12 +87,36 @@ func (s *DsnSelector) Report(node *registry.Node, cost time.Duration, err error)
 	return nil
 }
 
+// Options are ResolvableSelector options.
+type Options struct {
+	Extractor       ServiceNameExtractor
+	EnableParseAddr bool
+}
+
+// Option sets ResolvableSelector options.
+type Option func(*Options)
+
+// WithExtractor ...
+func WithExtractor(extractor ServiceNameExtractor) Option {
+	return func(o *Options) {
+		o.Extractor = extractor
+	}
+}
+
+// WithEnableParseAddr ...
+func WithEnableParseAddr(enable bool) Option {
+	return func(o *Options) {
+		o.EnableParseAddr = enable
+	}
+}
+
 // ResolvableSelector dsn-selector with address resolver
 type ResolvableSelector struct {
 	dsns                 map[string]*registry.Node
 	lock                 sync.RWMutex
 	resolverSelectorName string
 	extractor            ServiceNameExtractor
+	EnableParseAddr      bool
 }
 
 // ServiceNameExtractor extracts the part of the service name in the dsn, and return the starting position and length
@@ -98,6 +139,25 @@ func NewResolvableSelector(selectorName string, extractor ServiceNameExtractor)
 	}
 }
 
+// NewResolvableSelectorWithOpts ...
+func NewResolvableSelectorWithOpts(selectorName string, opt ...Option) selector.Selector {
+	opts := &Options{
+		Extractor:       &URIHostExtractor{},
+		EnableParseAddr: true,
+	}
+
+	for _, o := range opt {
+		o(opts)
+	}
+
+	return &ResolvableSelector{
+		dsns:                 make(map[string]*registry.Node),
+		resolverSelectorName: selectorName,
+		extractor:            opts.Extractor,
+		EnableParseAddr:      opts.EnableParseAddr,
+	}
+}
+
 // Select selects address from mongodb+polaris://user:passwd@poloars_name/db
 func (s *ResolvableSelector) Select(serviceName string, opt ...selector.Option) (*registry.Node, error) {
 	// resolve serviceName from dsn
@@ -165,6 +225,13 @@ func (s *ResolvableSelector) dsnRW(address, serviceName string, resolvedNode *re
 			"resolved": resolvedNode,
 		},
 	}
+
+	if s.EnableParseAddr {
+		node.ParseAddr = func(network, address string) net.Addr {
+			return unresolvedAddr{resolvedNode.Network, resolvedNode.Address}
+		}
+	}
+
 	s.dsns[address] = node
 	return node
 }
@@ -193,3 +260,18 @@ func (s *ResolvableSelector) Report(node *registry.Node, cost time.Duration, err
 
 	return resolver.Report(resolvedNode, cost, err)
 }
+
+type unresolvedAddr struct {
+	network string
+	address string
+}
+
+// Network returns the unresolved original network.
+func (a unresolvedAddr) Network() string {
+	return a.network
+}
+
+// String returns the unresolved original address.
+func (a unresolvedAddr) String() string {
+	return a.address
+}

dsn_test.go

@@ -20,7 +20,7 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-
+	"github.com/stretchr/testify/require"
 	"trpc.group/trpc-go/trpc-go/naming/registry"
 	"trpc.group/trpc-go/trpc-go/naming/selector"
 	dsn "trpc.group/trpc-go/trpc-selector-dsn"
@@ -141,3 +141,34 @@ func TestResolvableSelector(t *testing.T) {
 	_, err = noExtractorS.Select("123")
 	assert.EqualError(err, "service name extractor can not be nil", "case: empty")
 }
+
+func TestUnresolvedAddress(t *testing.T) {
+	n, err := dsn.NewDsnSelector(true).Select(t.Name())
+	require.Nil(t, err)
+	require.NotNil(t, n.ParseAddr)
+	addr := n.ParseAddr(n.Network, n.Address)
+	require.Equal(t, t.Name(), addr.String())
+	require.Equal(t, "", addr.Network())
+}
+
+func TestResolvableSelectorParseAddr(t *testing.T) {
+	target := "root:root1234@tcp(127.0.0.1:3306)/redis_db?timeout=1s"
+
+	// old selector
+	n, err := dsn.NewResolvableSelector("dsn", &dsn.URIHostExtractor{}).Select(target)
+	require.Nil(t, err)
+	require.Nil(t, n.ParseAddr)
+
+	// new selector with parseaddr
+	n, err = dsn.NewResolvableSelectorWithOpts("dsn").Select(target)
+	require.Nil(t, err)
+	require.NotNil(t, n.ParseAddr)
+	addr := n.ParseAddr(n.Network, n.Address)
+	require.Equal(t, "127.0.0.1:3306", addr.String())
+	require.Equal(t, "", addr.Network())
+
+	// new selector disable parseaddr
+	n, err = dsn.NewResolvableSelectorWithOpts("dsn", dsn.WithEnableParseAddr(false)).Select(target)
+	require.Nil(t, err)
+	require.Nil(t, n.ParseAddr)
+}

go.mod

@@ -4,11 +4,13 @@ go 1.18
 
 require (
 	github.com/stretchr/testify v1.8.0
-	trpc.group/trpc-go/trpc-go v1.0.0
+	trpc.group/trpc-go/trpc-go v1.0.3
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/kr/pretty v0.1.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,17 +1,23 @@
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
+github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
+github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
+github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
+gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
-trpc.group/trpc-go/trpc-go v1.0.0 h1:bSbcNpRFEXJONkwMVs8Xd+Mw/mFPUeE9Lcxk7KAaSoU=
-trpc.group/trpc-go/trpc-go v1.0.0/go.mod h1:ve2YyZleGVbnKr0RLUJcu35dXw2zZmsi3RdKVPgL4+4=
+trpc.group/trpc-go/trpc-go v1.0.3 h1:X4RhPmJOkVoK6EGKoV241dvEpB6EagBeyu3ZrqkYZQY=
+trpc.group/trpc-go/trpc-go v1.0.3/go.mod h1:82O+G2rD5ST+JAPuPPSqvsr6UI59UxV27iAILSkAIlQ=

service_name_extractor.go

@@ -36,45 +36,28 @@ func (e *URIHostExtractor) Extract(uri string) (int, int, error) {
 
 	//  resolve end-part of the host
 	begin := offset
-	length, err := dealHostEndPart(uri)
-	if err != nil {
-		return 0, 0, err
-	}
-	uri = uri[0:length]
-
-	return e.dealProtocolToken(uri, begin, length)
-}
-
-func dealHostEndPart(uri string) (int, error) {
 	length := len(uri)
 	if idx := strings.IndexAny(uri, "/?@"); idx != -1 {
 		if uri[idx] == '@' {
-			return 0, errors.New("parse host from uri: unescaped @ sign in user info")
+			return 0, 0, errors.New("parse host from uri: unescaped @ sign in user info")
 		}
 		if uri[idx] == '?' {
-			return 0, errors.New("parse host from uri: must have a / before the query ?")
+			return 0, 0, errors.New("parse host from uri: must have a / before the query ?")
 		}
 		length = idx
 	}
-	return length, nil
-}
+	uri = uri[0:length]
 
-func (e *URIHostExtractor) dealProtocolToken(uri string, begin, length int) (int, int, error) {
-	begin, length = dealProtocolPrefix(uri, begin, length)
-	length = dealProtocolSuffix(uri, length)
-	return begin, length, nil
+	return e.dealProtocolToken(uri, begin, length)
 }
 
-func dealProtocolPrefix(uri string, begin, length int) (int, int) {
+func (e *URIHostExtractor) dealProtocolToken(uri string, begin, length int) (int, int, error) {
 	if strings.HasPrefix(uri, "tcp(") {
-		return begin + 4, length - 4
+		begin += 4
+		length -= 4
 	}
-	return begin, length
-}
-
-func dealProtocolSuffix(uri string, length int) int {
 	if strings.HasSuffix(uri, ")") {
-		return length - 1
+		length--
 	}
-	return length
+	return begin, length, nil
 }