vespa-cli: Build endpoints map

2 files changed, 47 insertions, 16 deletions

diff --git a/client/go/vespa/target.go b/client/go/vespa/target.go
index f497bf5b3cd..41199b64137 100644
--- a/client/go/vespa/target.go
+++ b/client/go/vespa/target.go
@@ -208,11 +208,34 @@ type cloudTarget struct {
 	tlsOptions TLSOptions
 	logOptions LogOptions
 
-	queryURL       string
-	documentURL    string
+	urlsByCluster  map[string]string
 	authConfigPath string
 }
 
+func (t *cloudTarget) resolveEndpoint(cluster string) (string, error) {
+	if cluster == "" {
+		for _, u := range t.urlsByCluster {
+			if len(t.urlsByCluster) == 1 {
+				return u, nil
+			} else {
+				return "", fmt.Errorf("multiple clusters, none chosen: %v", t.urlsByCluster)
+			}
+		}
+	} else {
+		u := t.urlsByCluster[cluster]
+		if u == "" {
+			clusters := make([]string, len(t.urlsByCluster))
+			for c := range t.urlsByCluster {
+				clusters = append(clusters, c)
+			}
+			return "", fmt.Errorf("unknown cluster '%s': must be one of %v, ", cluster, clusters)
+		}
+		return u, nil
+	}
+
+	return "", fmt.Errorf("no endpoints")
+}
+
 func (t *cloudTarget) Type() string { return t.targetType }
 
 func (t *cloudTarget) Service(name string, timeout time.Duration, runID int64) (*Service, error) {
@@ -225,15 +248,17 @@ func (t *cloudTarget) Service(name string, timeout time.Duration, runID int64) (
 	case deployService:
 		return &Service{Name: name, BaseURL: t.apiURL}, nil
 	case queryService:
-		if t.queryURL == "" {
-			return nil, fmt.Errorf("service %s is not discovered", name)
+		queryURL, err := t.resolveEndpoint("")
+		if err != nil {
+			return nil, err
 		}
-		return &Service{Name: name, BaseURL: t.queryURL, TLSOptions: t.tlsOptions}, nil
+		return &Service{Name: name, BaseURL: queryURL, TLSOptions: t.tlsOptions}, nil
 	case documentService:
-		if t.documentURL == "" {
-			return nil, fmt.Errorf("service %s is not discovered", name)
+		documentURL, err := t.resolveEndpoint("")
+		if err != nil {
+			return nil, err
 		}
-		return &Service{Name: name, BaseURL: t.documentURL, TLSOptions: t.tlsOptions}, nil
+		return &Service{Name: name, BaseURL: documentURL, TLSOptions: t.tlsOptions}, nil
 	}
 	return nil, fmt.Errorf("unknown service: %s", name)
 }
@@ -402,7 +427,7 @@ func (t *cloudTarget) discoverEndpoints(timeout time.Duration) error {
 	if err := t.PrepareApiRequest(req, t.deployment.Application.SerializedForm()); err != nil {
 		return err
 	}
-	var endpointURL string
+	urlsByCluster := make(map[string]string)
 	endpointFunc := func(status int, response []byte) (bool, error) {
 		if ok, err := isOK(status); !ok {
 			return ok, err
@@ -414,17 +439,21 @@ func (t *cloudTarget) discoverEndpoints(timeout time.Duration) error {
 		if len(resp.Endpoints) == 0 {
 			return false, nil
 		}
-		endpointURL = resp.Endpoints[0].URL
+		for _, endpoint := range resp.Endpoints {
+			if endpoint.Scope != "zone" {
+				continue
+			}
+			urlsByCluster[endpoint.Cluster] = endpoint.URL
+		}
 		return true, nil
 	}
 	if _, err = wait(endpointFunc, func() *http.Request { return req }, &t.tlsOptions.KeyPair, timeout); err != nil {
 		return err
 	}
-	if endpointURL == "" {
-		return fmt.Errorf("no endpoint discovered")
+	if len(urlsByCluster) == 0 {
+		return fmt.Errorf("no endpoints discovered")
 	}
-	t.queryURL = endpointURL
-	t.documentURL = endpointURL
+	t.urlsByCluster = urlsByCluster
 	return nil
 }
 
@@ -460,7 +489,9 @@ func CloudTarget(apiURL string, deployment Deployment, apiKey []byte, tlsOptions
 }
 
 type deploymentEndpoint struct {
-	URL string `json:"url"`
+	Cluster string `json:"cluster"`
+	URL     string `json:"url"`
+	Scope   string `json:"scope"`
 }
 
 type deploymentResponse struct {
diff --git a/client/go/vespa/target_test.go b/client/go/vespa/target_test.go
index d4d23901513..1ed60efad0a 100644
--- a/client/go/vespa/target_test.go
+++ b/client/go/vespa/target_test.go
@@ -25,7 +25,7 @@ func (v *mockVespaApi) mockVespaHandler(w http.ResponseWriter, req *http.Request
 	case "/application/v4/tenant/t1/application/a1/instance/i1/environment/dev/region/us-north-1":
 		response := "{}"
 		if v.deploymentConverged {
-			response = fmt.Sprintf(`{"endpoints": [{"url": "%s"}]}`, v.serverURL)
+			response = fmt.Sprintf(`{"endpoints": [{"url": "%s","scope": "zone","cluster": "cluster1"}]}`, v.serverURL)
 		}
 		w.Write([]byte(response))
 	case "/application/v4/tenant/t1/application/a1/instance/i1/job/dev-us-north-1/run/42":