Support feeding to hosted target

6 files changed, 44 insertions, 27 deletions

diff --git a/client/go/internal/cli/cmd/feed.go b/client/go/internal/cli/cmd/feed.go
index 5b168ef79a2..2ea3ee1c4ed 100644
--- a/client/go/internal/cli/cmd/feed.go
+++ b/client/go/internal/cli/cmd/feed.go
@@ -93,15 +93,30 @@ $ cat docs.jsonl | vespa feed -`,
 	return cmd
 }
 
-func createServiceClients(service *vespa.Service, n int) []util.HTTPClient {
-	clients := make([]util.HTTPClient, 0, n)
+func createServices(n int, timeout time.Duration, cli *CLI) ([]util.HTTPClient, string, error) {
+	if n < 1 {
+		return nil, "", fmt.Errorf("need at least one client")
+	}
+	target, err := cli.target(targetOptions{})
+	if err != nil {
+		return nil, "", err
+	}
+	services := make([]util.HTTPClient, 0, n)
+	baseURL := ""
 	for i := 0; i < n; i++ {
-		client := service.Client().Clone()
+		service, err := cli.service(target, vespa.DocumentService, 0, cli.config.cluster())
+		if err != nil {
+			return nil, "", err
+		}
+		baseURL = service.BaseURL
+		// Create a separate HTTP client for each service
+		client := cli.httpClientFactory(timeout)
 		// Feeding should always use HTTP/2
 		util.ForceHTTP2(client, service.TLSOptions.KeyPair, service.TLSOptions.CACertificate, service.TLSOptions.TrustAll)
-		clients = append(clients, client)
+		service.SetClient(client)
+		services = append(services, service)
 	}
-	return clients
+	return services, baseURL, nil
 }
 
 func summaryTicker(secs int, cli *CLI, start time.Time, statsFunc func() document.Stats) *time.Ticker {
@@ -130,21 +145,21 @@ func (opts feedOptions) compressionMode() (document.Compression, error) {
 }
 
 func feed(files []string, options feedOptions, cli *CLI) error {
-	service, err := documentService(cli)
+	timeout := time.Duration(options.timeoutSecs) * time.Second
+	clients, baseURL, err := createServices(options.connections, timeout, cli)
 	if err != nil {
 		return err
 	}
-	clients := createServiceClients(service, options.connections)
 	compression, err := options.compressionMode()
 	if err != nil {
 		return err
 	}
 	client, err := document.NewClient(document.ClientOptions{
 		Compression: compression,
-		Timeout:     time.Duration(options.timeoutSecs) * time.Second,
+		Timeout:     timeout,
 		Route:       options.route,
 		TraceLevel:  options.traceLevel,
-		BaseURL:     service.BaseURL,
+		BaseURL:     baseURL,
 		NowFunc:     cli.now,
 	}, clients)
 	if err != nil {
diff --git a/client/go/internal/cli/cmd/feed_test.go b/client/go/internal/cli/cmd/feed_test.go
index 467d55a0a6e..097d4ae5fa3 100644
--- a/client/go/internal/cli/cmd/feed_test.go
+++ b/client/go/internal/cli/cmd/feed_test.go
@@ -24,10 +24,9 @@ func (c *manualClock) now() time.Time {
 }
 
 func TestFeed(t *testing.T) {
-	httpClient := &mock.HTTPClient{}
 	clock := &manualClock{tick: time.Second}
 	cli, stdout, stderr := newTestCLI(t)
-	cli.httpClient = httpClient
+	httpClient := cli.httpClient.(*mock.HTTPClient)
 	cli.now = clock.now
 
 	td := t.TempDir()
diff --git a/client/go/internal/cli/cmd/root.go b/client/go/internal/cli/cmd/root.go
index 1b37ff00269..c4012024426 100644
--- a/client/go/internal/cli/cmd/root.go
+++ b/client/go/internal/cli/cmd/root.go
@@ -47,13 +47,14 @@ type CLI struct {
 	config  *Config
 	version version.Version
 
-	httpClient   util.HTTPClient
-	auth0Factory auth0Factory
-	ztsFactory   ztsFactory
-	exec         executor
-	isTerminal   func() bool
-	spinner      func(w io.Writer, message string, fn func() error) error
-	now          func() time.Time
+	httpClient        util.HTTPClient
+	httpClientFactory func(timeout time.Duration) util.HTTPClient
+	auth0Factory      auth0Factory
+	ztsFactory        ztsFactory
+	exec              executor
+	isTerminal        func() bool
+	spinner           func(w io.Writer, message string, fn func() error) error
+	now               func() time.Time
 }
 
 // ErrCLI is an error returned to the user. It wraps an exit status, a regular error and optional hints for resolving
@@ -122,17 +123,19 @@ For detailed description of flags and configuration, see 'vespa help config'.
 	if err != nil {
 		return nil, err
 	}
+	httpClientFactory := util.CreateClient
 	cli := CLI{
 		Environment: env,
 		Stdin:       os.Stdin,
 		Stdout:      stdout,
 		Stderr:      stderr,
 
-		version:    version,
-		cmd:        cmd,
-		httpClient: util.CreateClient(time.Second * 10),
-		exec:       &execSubprocess{},
-		now:        time.Now,
+		version:           version,
+		cmd:               cmd,
+		httpClient:        httpClientFactory(time.Second * 10),
+		httpClientFactory: httpClientFactory,
+		exec:              &execSubprocess{},
+		now:               time.Now,
 		auth0Factory: func(httpClient util.HTTPClient, options auth0.Options) (vespa.Authenticator, error) {
 			return auth0.NewClient(httpClient, options)
 		},
diff --git a/client/go/internal/cli/cmd/testutil_test.go b/client/go/internal/cli/cmd/testutil_test.go
index 492e40d8855..61d6c15c5a0 100644
--- a/client/go/internal/cli/cmd/testutil_test.go
+++ b/client/go/internal/cli/cmd/testutil_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/vespa-engine/vespa/client/go/internal/cli/auth/auth0"
 	"github.com/vespa-engine/vespa/client/go/internal/mock"
@@ -28,6 +29,7 @@ func newTestCLI(t *testing.T, envVars ...string) (*CLI, *bytes.Buffer, *bytes.Bu
 		t.Fatal(err)
 	}
 	httpClient := &mock.HTTPClient{}
+	cli.httpClientFactory = func(timeout time.Duration) util.HTTPClient { return httpClient }
 	cli.httpClient = httpClient
 	cli.exec = &mock.Exec{}
 	cli.auth0Factory = func(httpClient util.HTTPClient, options auth0.Options) (vespa.Authenticator, error) {
diff --git a/client/go/internal/util/http.go b/client/go/internal/util/http.go
index 8a67b24dffb..e49d1254f26 100644
--- a/client/go/internal/util/http.go
+++ b/client/go/internal/util/http.go
@@ -16,7 +16,6 @@ import (
 
 type HTTPClient interface {
 	Do(request *http.Request, timeout time.Duration) (response *http.Response, error error)
-	Clone() HTTPClient
 }
 
 type defaultHTTPClient struct {
@@ -34,8 +33,6 @@ func (c *defaultHTTPClient) Do(request *http.Request, timeout time.Duration) (re
 	return c.client.Do(request)
 }
 
-func (c *defaultHTTPClient) Clone() HTTPClient { return CreateClient(c.client.Timeout) }
-
 func ConfigureTLS(client HTTPClient, certificates []tls.Certificate, caCertificate []byte, trustAll bool) {
 	c, ok := client.(*defaultHTTPClient)
 	if !ok {
diff --git a/client/go/internal/vespa/target.go b/client/go/internal/vespa/target.go
index 9f3fd7f5c65..6dd64dd1275 100644
--- a/client/go/internal/vespa/target.go
+++ b/client/go/internal/vespa/target.go
@@ -110,7 +110,8 @@ func (s *Service) Do(request *http.Request, timeout time.Duration) (*http.Respon
 	return s.httpClient.Do(request, timeout)
 }
 
-func (s *Service) Client() util.HTTPClient { return s.httpClient }
+// SetClient sets the HTTP client that this service should use.
+func (s *Service) SetClient(client util.HTTPClient) { s.httpClient = client }
 
 // Wait polls the health check of this service until it succeeds or timeout passes.
 func (s *Service) Wait(timeout time.Duration) (int, error) {