Send the Vizier ID when requesting Vizier configs from the cloud

src/operator/controllers/BUILD.bazel

@@ -29,15 +29,18 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//src/api/proto/cloudpb:cloudapi_pl_go_proto",
+        "//src/api/proto/uuidpb:uuid_pl_go_proto",
         "//src/api/proto/vizierconfigpb:vizier_pl_go_proto",
         "//src/operator/apis/px.dev/v1alpha1",
         "//src/shared/goversion",
         "//src/shared/services",
         "//src/shared/status",
+        "//src/utils",
         "//src/utils/shared/certs",
         "//src/utils/shared/k8s",
         "@com_github_blang_semver//:semver",
         "@com_github_cenkalti_backoff_v4//:backoff",
+        "@com_github_gofrs_uuid//:uuid",
         "@com_github_gogo_protobuf//types",
         "@com_github_sirupsen_logrus//:logrus",
         "@io_k8s_api//apps/v1:apps",
@@ -64,13 +67,16 @@ pl_go_test(
         "monitor_test.go",
         "node_watcher_test.go",
         "pvc_watcher_test.go",
+        "vizier_controller_test.go",
     ],
     embed = [":controllers"],
     deps = [
         "//src/api/proto/cloudpb:cloudapi_pl_go_proto",
         "//src/api/proto/cloudpb/mock",
+        "//src/api/proto/uuidpb:uuid_pl_go_proto",
         "//src/operator/apis/px.dev/v1alpha1",
         "//src/shared/status",
+        "//src/utils",
         "@com_github_gogo_protobuf//types",
         "@com_github_golang_mock//gomock",
         "@com_github_stretchr_testify//assert",
@@ -83,5 +89,6 @@ pl_go_test(
         "@io_k8s_client_go//kubernetes/fake",
         "@io_k8s_client_go//testing",
         "@io_k8s_sigs_controller_runtime//pkg/client",
+        "@org_golang_google_grpc//:grpc",
     ],
 )

src/operator/controllers/vizier_controller.go

@@ -32,6 +32,7 @@ import (
 
 	"github.com/blang/semver"
 	"github.com/cenkalti/backoff/v4"
+	"github.com/gofrs/uuid"
 	"github.com/gogo/protobuf/types"
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
@@ -48,11 +49,13 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	"px.dev/pixie/src/api/proto/cloudpb"
+	"px.dev/pixie/src/api/proto/uuidpb"
 	"px.dev/pixie/src/api/proto/vizierconfigpb"
 	"px.dev/pixie/src/operator/apis/px.dev/v1alpha1"
 	version "px.dev/pixie/src/shared/goversion"
 	"px.dev/pixie/src/shared/services"
 	"px.dev/pixie/src/shared/status"
+	"px.dev/pixie/src/utils"
 	"px.dev/pixie/src/utils/shared/certs"
 	"px.dev/pixie/src/utils/shared/k8s"
 )
@@ -497,7 +500,12 @@ func (r *VizierReconciler) deployVizier(ctx context.Context, req ctrl.Request, v
 		return err
 	}
 
-	configForVizierResp, err := generateVizierYAMLsConfig(ctx, req.Namespace, r.K8sVersion, vz, cloudClient)
+	// The Vizier ID is only available once the cloud connector has registered the Vizier, so it is
+	// nil on the initial deploy. Pixie Cloud uses it to look up the Vizier's org when the deploy
+	// key is no longer valid, for example during an update.
+	vizierID := getVizierID(r.Clientset, req.Namespace)
+
+	configForVizierResp, err := generateVizierYAMLsConfig(ctx, req.Namespace, r.K8sVersion, vizierID, vz, cloudpb.NewConfigServiceClient(cloudClient))
 	if err != nil {
 		log.WithError(err).Error("Failed to generate configs for Vizier YAMLs")
 		return err
@@ -809,14 +817,13 @@ func convertResourceType(originalLst v1.ResourceList) *vizierconfigpb.ResourceLi
 
 // generateVizierYAMLsConfig is responsible retrieving a yaml map of configurations from
 // Pixie Cloud.
-func generateVizierYAMLsConfig(ctx context.Context, ns string, k8sVersion string, vz *v1alpha1.Vizier, conn *grpc.ClientConn) (*cloudpb.ConfigForVizierResponse,
+func generateVizierYAMLsConfig(ctx context.Context, ns string, k8sVersion string, vizierID *uuidpb.UUID, vz *v1alpha1.Vizier, client cloudpb.ConfigServiceClient) (*cloudpb.ConfigForVizierResponse,
 	error,
 ) {
-	client := cloudpb.NewConfigServiceClient(conn)
-
 	req := &cloudpb.ConfigForVizierRequest{
 		Namespace:  ns,
 		K8sVersion: k8sVersion,
+		VizierID:   vizierID,
 		VzSpec: &vizierconfigpb.VizierSpec{
 			Version:               vz.Spec.Version,
 			DeployKey:             vz.Spec.DeployKey,
@@ -1036,6 +1043,20 @@ func updatePodSpec(nodeSelector map[string]string, tolerations []v1.Toleration,
 	podSpec["securityContext"] = sCtx
 }
 
+// getVizierID returns the ID of the Vizier stored in the pl-cluster-secrets secret, or nil if
+// the Vizier has not been registered with Pixie Cloud yet.
+func getVizierID(clientset kubernetes.Interface, namespace string) *uuidpb.UUID {
+	s := k8s.GetSecret(clientset, namespace, "pl-cluster-secrets")
+	if s == nil {
+		return nil
+	}
+	id, err := uuid.FromString(string(s.Data["cluster-id"]))
+	if err != nil {
+		return nil
+	}
+	return utils.ProtoFromUUID(id)
+}
+
 func waitForCluster(clientset *kubernetes.Clientset, namespace string) error {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
 	defer cancel()

src/operator/controllers/vizier_controller_test.go

@@ -0,0 +1,136 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package controllers
+
+import (
+	"context"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/kubernetes/fake"
+
+	"px.dev/pixie/src/api/proto/cloudpb"
+	mock_cloudpb "px.dev/pixie/src/api/proto/cloudpb/mock"
+	"px.dev/pixie/src/api/proto/uuidpb"
+	"px.dev/pixie/src/operator/apis/px.dev/v1alpha1"
+	"px.dev/pixie/src/utils"
+)
+
+const testVizierID = "7ba7b810-9dad-11d1-80b4-00c04fd430c8"
+
+func TestGetVizierID(t *testing.T) {
+	tests := []struct {
+		name       string
+		secretData map[string][]byte
+		expectedID *uuidpb.UUID
+	}{
+		{
+			name:       "missing secret",
+			secretData: nil,
+			expectedID: nil,
+		},
+		{
+			name:       "missing cluster ID",
+			secretData: map[string][]byte{},
+			expectedID: nil,
+		},
+		{
+			name: "invalid cluster ID",
+			secretData: map[string][]byte{
+				"cluster-id": []byte("not-a-uuid"),
+			},
+			expectedID: nil,
+		},
+		{
+			name: "valid cluster ID",
+			secretData: map[string][]byte{
+				"cluster-id": []byte(testVizierID),
+			},
+			expectedID: utils.ProtoFromUUIDStrOrNil(testVizierID),
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var objects []runtime.Object
+			if test.secretData != nil {
+				objects = append(objects, &v1.Secret{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "pl-cluster-secrets",
+						Namespace: "pl",
+					},
+					Data: test.secretData,
+				})
+			}
+			clientset := fake.NewSimpleClientset(objects...)
+
+			assert.Equal(t, test.expectedID, getVizierID(clientset, "pl"))
+		})
+	}
+}
+
+func TestGenerateVizierYAMLsConfig_VizierID(t *testing.T) {
+	tests := []struct {
+		name     string
+		vizierID *uuidpb.UUID
+	}{
+		{
+			name:     "initial deploy",
+			vizierID: nil,
+		},
+		{
+			name:     "update",
+			vizierID: utils.ProtoFromUUIDStrOrNil(testVizierID),
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			client := mock_cloudpb.NewMockConfigServiceClient(ctrl)
+
+			resp := &cloudpb.ConfigForVizierResponse{
+				NameToYamlContent: map[string]string{"vizier": "yaml"},
+			}
+			client.EXPECT().GetConfigForVizier(gomock.Any(), gomock.Any()).
+				DoAndReturn(func(_ context.Context, req *cloudpb.ConfigForVizierRequest, _ ...grpc.CallOption) (*cloudpb.ConfigForVizierResponse, error) {
+					assert.Equal(t, "pl", req.Namespace)
+					assert.Equal(t, "1.30.0", req.K8sVersion)
+					assert.Equal(t, test.vizierID, req.VizierID)
+					assert.Equal(t, "deploy-key", req.VzSpec.DeployKey)
+					return resp, nil
+				})
+
+			vz := &v1alpha1.Vizier{
+				Spec: v1alpha1.VizierSpec{
+					DeployKey: "deploy-key",
+					Pod:       &v1alpha1.PodPolicy{},
+				},
+			}
+			actualResp, err := generateVizierYAMLsConfig(context.Background(), "pl", "1.30.0", test.vizierID, vz, client)
+			assert.NoError(t, err)
+			assert.Equal(t, resp, actualResp)
+		})
+	}
+}