Allow VizierID to be used to find the Vizier's org (#1650)

Summary: This PR allows the VizierID to be used when getting the org as
opposed to only relying on the DeployKey when adding feature flags.

Relevant Issues: #1632

Type of change: /kind bug

Test Plan: Another PR needs to come in on the operator side to
retrieve/send the VizierID to the cloud. Getting this PR in on the cloud
side will simplify the process of testing the integration between the
two later.

---------

Signed-off-by: Kartik Pattaswamy <kpattaswamy@pixielabs.ai>

src/api/proto/cloudpb/cloudapi.proto

@@ -956,6 +956,8 @@ message ConfigForVizierRequest {
   vizierconfigpb.VizierSpec vz_spec = 2;
   // Kubernetes version of the cluster Vizier is running on.
   string k8s_version = 3 [ (gogoproto.customname) = "K8sVersion" ];
+  // ID of the Vizier (Cluster ID).
+  string vizier_id = 4 [ (gogoproto.customname) = "VizierID" ];
 }
 
 // ConfigForVizierResponse is the response to a ConfigForVizierRequest.

src/cloud/api/controllers/config_grpc.go

@@ -62,6 +62,7 @@ func (c *ConfigServiceServer) GetConfigForVizier(ctx context.Context,
 			Registry:                 vizSpecReq.Registry,
 		},
 		K8sVersion: req.K8sVersion,
+		VizierID:   req.VizierID,
 	})
 	if err != nil {
 		return nil, err

src/cloud/config_manager/config_manager_server.go

@@ -73,6 +73,20 @@ func newDeploymentKeyClient() (vzmgrpb.VZDeploymentKeyServiceClient, error) {
 	return vzmgrpb.NewVZDeploymentKeyServiceClient(deployKeyChannel), nil
 }
 
+func newVizierManagerClient() (vzmgrpb.VZMgrServiceClient, error) {
+	dialOpts, err := services.GetGRPCClientDialOpts()
+	if err != nil {
+		return nil, err
+	}
+
+	vizierManagerChannel, err := grpc.Dial(viper.GetString("vzmgr_service"), dialOpts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return vzmgrpb.NewVZMgrServiceClient(vizierManagerChannel), nil
+}
+
 func main() {
 	services.SetupService("config-manager-service", 50500)
 	services.PostFlagSetupAndParse()
@@ -93,7 +107,13 @@ func main() {
 	if err != nil {
 		log.WithError(err).Fatal("Could not connect with Artifact Service.")
 	}
-	svr := controllers.NewServer(atClient, deployKeyClient, viper.GetString("ld_sdk_key"))
+
+	vzmgrClient, err := newVizierManagerClient()
+	if err != nil {
+		log.WithError(err).Fatal("Could not connect with VizierManager Service.")
+	}
+
+	svr := controllers.NewServer(atClient, deployKeyClient, viper.GetString("ld_sdk_key"), vzmgrClient)
 	serverOpts := &server.GRPCServerOptions{
 		DisableAuth: map[string]bool{
 			"/px.services.ConfigManagerService/GetConfigForVizier":   true,

src/cloud/config_manager/configmanagerpb/service.proto

@@ -42,6 +42,8 @@ message ConfigForVizierRequest {
   vizierconfigpb.VizierSpec vz_spec = 2;
   // Kubernetes version of the cluster Vizier is running on.
   string k8s_version = 3 [ (gogoproto.customname) = "K8sVersion" ];
+  // ID of the Vizier (Cluster ID).
+  string vizier_id = 4 [ (gogoproto.customname) = "VizierID" ];
 }
 
 // ConfigForVizierResponse is the response to a ConfigForVizierRequest.

src/cloud/config_manager/controllers/server.go

@@ -52,14 +52,16 @@ type Server struct {
 	atClient            atpb.ArtifactTrackerClient
 	deployKeyClient     vzmgrpb.VZDeploymentKeyServiceClient
 	vzFeatureFlagClient VizierFeatureFlagClient
+	vzmgrClient         vzmgrpb.VZMgrServiceClient
 }
 
 // NewServer creates GRPC handlers.
-func NewServer(atClient atpb.ArtifactTrackerClient, deployKeyClient vzmgrpb.VZDeploymentKeyServiceClient, ldSDKKey string) *Server {
+func NewServer(atClient atpb.ArtifactTrackerClient, deployKeyClient vzmgrpb.VZDeploymentKeyServiceClient, ldSDKKey string, vzmgrClient vzmgrpb.VZMgrServiceClient) *Server {
 	return &Server{
 		atClient:            atClient,
 		deployKeyClient:     deployKeyClient,
 		vzFeatureFlagClient: NewVizierFeatureFlagClient(ldSDKKey),
+		vzmgrClient:         vzmgrClient,
 	}
 }
 
@@ -217,12 +219,25 @@ func (s *Server) GetConfigForVizier(ctx context.Context,
 	}
 	AddDefaultTableStoreSize(tmplValues.PEMMemoryRequest, tmplValues.CustomPEMFlags)
 
-	// Next we inject any feature flags that we want to set for this org.
-	orgID, err := s.getOrgIDForDeployKey(tmplValues.DeployKey)
+	// Attempt to get the org ID from DeployKey, otherwise from the Vizier.
+	var orgID uuid.UUID
+	orgID, err = s.getOrgIDForDeployKey(tmplValues.DeployKey)
 	if err != nil || orgID == uuid.Nil {
-		log.WithError(err).Error("Error getting org ID from deploy key, skipping feature flag logic")
-	} else {
+		log.WithError(err).Error("Error getting org ID from deploy key")
+	}
+	if orgID == uuid.Nil && in.VizierID != "" {
+		resp, err := s.vzmgrClient.GetOrgFromVizier(ctx, utils.ProtoFromUUIDStrOrNil(in.VizierID))
+		orgID = utils.UUIDFromProtoOrNil(resp.OrgID)
+		if err != nil || orgID == uuid.Nil {
+			log.WithError(err).Error("Error getting org ID from Vizier")
+		}
+	}
+
+	// Next we inject any feature flags that we want to set for this org.
+	if orgID != uuid.Nil {
 		AddFeatureFlagsToTemplate(s.vzFeatureFlagClient, orgID, tmplValues)
+	} else {
+		log.Error("Skipping feature flag logic")
 	}
 
 	vzYamls, err := yamls.ExecuteTemplatedYAMLs(templatedYAMLs, vizieryamls.VizierTmplValuesToArgs(tmplValues))