[vtadmin] custom discovery resolver

go/vt/vtadmin/cluster/cluster_test.go

@@ -36,6 +36,7 @@ import (
 	"vitess.io/vitess/go/vt/vitessdriver"
 	"vitess.io/vitess/go/vt/vtadmin/cluster"
 	"vitess.io/vitess/go/vt/vtadmin/cluster/discovery/fakediscovery"
+	"vitess.io/vitess/go/vt/vtadmin/cluster/resolver"
 	vtadminerrors "vitess.io/vitess/go/vt/vtadmin/errors"
 	"vitess.io/vitess/go/vt/vtadmin/testutil"
 	"vitess.io/vitess/go/vt/vtadmin/vtctldclient/fakevtctldclient"
@@ -2652,7 +2653,9 @@ func TestGetTablets(t *testing.T) {
 			Id:   "c1",
 			Name: "one",
 		},
-		Discovery: disco,
+		ResolverOptions: &resolver.Options{
+			Discovery: disco,
+		},
 	})
 	db.DialFunc = func(cfg vitessdriver.Configuration) (*sql.DB, error) {
 		return nil, assert.AnError

go/vt/vtadmin/cluster/discovery/discovery.go

@@ -53,6 +53,10 @@ type Discovery interface {
 	// return an address is not specified by the interface, and can be
 	// implementation-specific.
 	DiscoverVTGateAddr(ctx context.Context, tags []string) (string, error)
+	// DiscoverVTGateAddrs returns a list of addresses of vtgates found in the
+	// discovery service. This is semantically equivalent to the result of
+	// DiscoverVTGateAddr for each gate returned by a call to DiscoverVTGates.
+	DiscoverVTGateAddrs(ctx context.Context, tags []string) ([]string, error)
 	// DiscoverVTGates returns a list of vtgates found in the discovery service.
 	// Tags can optionally be used to filter gates. Order of the gates is not
 	// specified by the interface, and can be implementation-specific.
@@ -68,6 +72,10 @@ type Discovery interface {
 	// return an address is not specified by the interface, and can be
 	// implementation-specific.
 	DiscoverVtctldAddr(ctx context.Context, tags []string) (string, error)
+	// DiscoverVtctldAddrs returns a list of addresses of vtctlds found in the
+	// discovery service. This is semantically equivalent to the result of
+	// DiscoverVtctldAddr for each gate returned by a call to DiscoverVtctlds.
+	DiscoverVtctldAddrs(ctx context.Context, tags []string) ([]string, error)
 	// DiscoverVtctlds returns a list of vtctlds found in the discovery service.
 	// Tags can optionally be used to filter vtctlds. Order of the vtctlds is
 	// not specified by the interface, and can be implementation-specific.

go/vt/vtadmin/cluster/discovery/discovery_consul.go

@@ -228,6 +228,31 @@ func (c *ConsulDiscovery) DiscoverVTGateAddr(ctx context.Context, tags []string)
 	return addr, nil
 }
 
+// DiscoverVTGateAddrs is part of the Discovery interface.
+func (c *ConsulDiscovery) DiscoverVTGateAddrs(ctx context.Context, tags []string) ([]string, error) {
+	span, ctx := trace.NewSpan(ctx, "ConsulDiscovery.DiscoverVTGateAddrs")
+	defer span.Finish()
+
+	executeFQDNTemplate := false
+
+	vtgates, err := c.discoverVTGates(ctx, tags, executeFQDNTemplate)
+	if err != nil {
+		return nil, err
+	}
+
+	addrs := make([]string, len(vtgates))
+	for i, vtgate := range vtgates {
+		addr, err := textutil.ExecuteTemplate(c.vtgateAddrTmpl, vtgate)
+		if err != nil {
+			return nil, fmt.Errorf("failed to execute vtgate address template for %v: %w", vtgate, err)
+		}
+
+		addrs[i] = addr
+	}
+
+	return addrs, nil
+}
+
 // DiscoverVTGates is part of the Discovery interface.
 func (c *ConsulDiscovery) DiscoverVTGates(ctx context.Context, tags []string) ([]*vtadminpb.VTGate, error) {
 	span, ctx := trace.NewSpan(ctx, "ConsulDiscovery.DiscoverVTGates")
@@ -348,6 +373,31 @@ func (c *ConsulDiscovery) DiscoverVtctldAddr(ctx context.Context, tags []string)
 	return addr, nil
 }
 
+// DiscoverVtctldAddrs is part of the Discovery interface.
+func (c *ConsulDiscovery) DiscoverVtctldAddrs(ctx context.Context, tags []string) ([]string, error) {
+	span, ctx := trace.NewSpan(ctx, "ConsulDiscovery.DiscoverVtctldAddrs")
+	defer span.Finish()
+
+	executeFQDNTemplate := false
+
+	vtctlds, err := c.discoverVtctlds(ctx, tags, executeFQDNTemplate)
+	if err != nil {
+		return nil, err
+	}
+
+	addrs := make([]string, len(vtctlds))
+	for i, vtctld := range vtctlds {
+		addr, err := textutil.ExecuteTemplate(c.vtctldAddrTmpl, vtctld)
+		if err != nil {
+			return nil, fmt.Errorf("failed to execute vtctld address template for %v: %w", vtctld, err)
+		}
+
+		addrs[i] = addr
+	}
+
+	return addrs, nil
+}
+
 // DiscoverVtctlds is part of the Discovery interface.
 func (c *ConsulDiscovery) DiscoverVtctlds(ctx context.Context, tags []string) ([]*vtadminpb.Vtctld, error) {
 	span, ctx := trace.NewSpan(ctx, "ConsulDiscovery.DiscoverVtctlds")

go/vt/vtadmin/cluster/discovery/discovery_json.go

@@ -23,6 +23,7 @@ import (
 	"math/rand"
 
 	"vitess.io/vitess/go/trace"
+
 	vtadminpb "vitess.io/vitess/go/vt/proto/vtadmin"
 )
 
@@ -146,6 +147,24 @@ func (d *JSONDiscovery) DiscoverVTGateAddr(ctx context.Context, tags []string) (
 	return gate.Hostname, nil
 }
 
+// DiscoverVTGateAddrs is part of the Discovery interface.
+func (d *JSONDiscovery) DiscoverVTGateAddrs(ctx context.Context, tags []string) ([]string, error) {
+	span, ctx := trace.NewSpan(ctx, "JSONDiscovery.DiscoverVTGateAddrs")
+	defer span.Finish()
+
+	gates, err := d.discoverVTGates(ctx, tags)
+	if err != nil {
+		return nil, err
+	}
+
+	addrs := make([]string, len(gates))
+	for i, gate := range gates {
+		addrs[i] = gate.Hostname
+	}
+
+	return addrs, nil
+}
+
 // DiscoverVTGates is part of the Discovery interface.
 func (d *JSONDiscovery) DiscoverVTGates(ctx context.Context, tags []string) ([]*vtadminpb.VTGate, error) {
 	span, ctx := trace.NewSpan(ctx, "JSONDiscovery.DiscoverVTGates")
@@ -228,6 +247,24 @@ func (d *JSONDiscovery) DiscoverVtctldAddr(ctx context.Context, tags []string) (
 	return vtctld.Hostname, nil
 }
 
+// DiscoverVtctldAddrs is part of the Discovery interface.
+func (d *JSONDiscovery) DiscoverVtctldAddrs(ctx context.Context, tags []string) ([]string, error) {
+	span, ctx := trace.NewSpan(ctx, "JSONDiscovery.DiscoverVtctldAddrs")
+	defer span.Finish()
+
+	vtctlds, err := d.discoverVtctlds(ctx, tags)
+	if err != nil {
+		return nil, err
+	}
+
+	addrs := make([]string, len(vtctlds))
+	for i, vtctld := range vtctlds {
+		addrs[i] = vtctld.Hostname
+	}
+
+	return addrs, nil
+}
+
 // DiscoverVtctlds is part of the Discovery interface.
 func (d *JSONDiscovery) DiscoverVtctlds(ctx context.Context, tags []string) ([]*vtadminpb.Vtctld, error) {
 	span, ctx := trace.NewSpan(ctx, "JSONDiscovery.DiscoverVtctlds")

go/vt/vtadmin/cluster/discovery/fakediscovery/discovery.go

@@ -181,6 +181,21 @@ func (d *Fake) DiscoverVTGateAddr(ctx context.Context, tags []string) (string, e
 	return gate.Hostname, nil
 }
 
+// DiscoverVTGateAddrs is part of the discovery.Discovery interface.
+func (d *Fake) DiscoverVTGateAddrs(ctx context.Context, tags []string) ([]string, error) {
+	gates, err := d.DiscoverVTGates(ctx, tags)
+	if err != nil {
+		return nil, err
+	}
+
+	addrs := make([]string, len(gates))
+	for i, gate := range gates {
+		addrs[i] = gate.Hostname
+	}
+
+	return addrs, nil
+}
+
 // DiscoverVtctlds is part of the discover.Discovery interface.
 func (d *Fake) DiscoverVtctlds(ctx context.Context, tags []string) ([]*vtadminpb.Vtctld, error) {
 	if d.vtctlds.shouldErr {
@@ -234,6 +249,21 @@ func (d *Fake) DiscoverVtctldAddr(ctx context.Context, tags []string) (string, e
 	return vtctld.Hostname, nil
 }
 
+// DiscoverVtctldAddrs is part of the discovery.Discovery interface.
+func (d *Fake) DiscoverVtctldAddrs(ctx context.Context, tags []string) ([]string, error) {
+	vtctlds, err := d.DiscoverVtctlds(ctx, tags)
+	if err != nil {
+		return nil, err
+	}
+
+	addrs := make([]string, len(vtctlds))
+	for i, vtctld := range vtctlds {
+		addrs[i] = vtctld.Hostname
+	}
+
+	return addrs, nil
+}
+
 // DiscoverVtctld is part of the discover.Discovery interface.
 func (d *Fake) DiscoverVtctld(ctx context.Context, tags []string) (*vtadminpb.Vtctld, error) {
 	vtctlds, err := d.DiscoverVtctlds(ctx, tags)