Allow VizierID to be used to find the Vizier's org

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/api/controllers/config_grpc_test.go

@@ -77,6 +77,7 @@ func TestConfigForVizier(t *testing.T) {
  Namespace: "test-namespace",
  VzSpec: vzSpec,
  K8sVersion: "1.24.1",
+ VizierID: "test-id",
  }
 
  nameToYamlContent := make(map[string]string)
@@ -92,6 +93,7 @@ func TestConfigForVizier(t *testing.T) {
  Namespace: "test-namespace",
  VzSpec: vzSpec,
  K8sVersion: "1.24.1",
+ VizierID: "test-id",
  })
  require.NoError(t, err)
  assert.Equal(t, resp.NameToYamlContent["fileAName"], "fileAContent")

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))