Provisioning a cluster should return clusterName along with clusterID

Summary:
We would like to know the name assigned to a cluster on the vizier side.
So when provisioning a cluster, return clusterName along with the clusterID.

Test Plan: Added unit tests that verify the field gets set as expected.

Reviewers: michelle, philkuz, nserrino

Reviewed By: philkuz

Signed-off-by: Vihang Mehta <vihang@pixielabs.ai>

Differential Revision: https://phab.corp.pixielabs.ai/D11029

GitOrigin-RevId: bdc859f

src/cloud/vzconn/bridge/grpc.go

@@ -91,7 +91,8 @@ func (s *GRPCServer) RegisterVizierDeployment(ctx context.Context, req *vzconnpb
 	}
 
 	return &vzconnpb.RegisterVizierDeploymentResponse{
-		VizierID: vzmgrResp.VizierID,
+		VizierID:   vzmgrResp.VizierID,
+		VizierName: vzmgrResp.VizierName,
 	}, nil
 }
 

src/cloud/vzconn/bridge/grpc_test.go

@@ -373,7 +373,7 @@ func TestNATSGRPCBridge_RegisterVizierDeployment(t *testing.T) {
 			DeploymentKey:  "deploy-key",
 			K8sClusterName: "some name",
 		}).
-		Return(&vzmgrpb.RegisterVizierDeploymentResponse{VizierID: utils.ProtoFromUUID(vizierID)}, nil)
+		Return(&vzmgrpb.RegisterVizierDeploymentResponse{VizierID: utils.ProtoFromUUID(vizierID), VizierName: "some_name"}, nil)
 
 	// Make some GRPC Requests.
 	ctx := context.Background()

src/cloud/vzconn/vzconnpb/service.proto

@@ -68,4 +68,6 @@ message RegisterVizierDeploymentRequest {
 message RegisterVizierDeploymentResponse {
   // The ID for this Vizier.
   uuidpb.UUID vizier_id = 1 [(gogoproto.customname) = "VizierID"];
+  // The name for this Vizier.
+  string vizier_name = 2;
 }

src/cloud/vzmgr/controllers/server.go

@@ -957,7 +957,7 @@ func findVizierWithEmptyUID(ctx context.Context, tx *sqlx.Tx, orgID uuid.UUID) (
 	return uuid.Nil, vizierStatus(cvmsgspb.VZ_ST_UNKNOWN), nil
 }
 
-func setClusterName(ctx context.Context, tx *sqlx.Tx, clusterID uuid.UUID, generateName func(i int) string) error {
+func setClusterName(ctx context.Context, tx *sqlx.Tx, clusterID uuid.UUID, generateName func(i int) string) (string, error) {
 	// Retry a few times until we find a name that doesn't collide.
 	finalName := ""
 	for rc := 0; rc < 10; rc++ {
@@ -973,21 +973,21 @@ func setClusterName(ctx context.Context, tx *sqlx.Tx, clusterID uuid.UUID, gener
 	}
 
 	if finalName == "" {
-		return errors.New("Could not find a unique cluster name")
+		return "", errors.New("Could not find a unique cluster name")
 	}
 
 	query := `UPDATE vizier_cluster SET cluster_name=$1 WHERE id=$2`
 	_, err := tx.ExecContext(ctx, query, finalName, clusterID)
 
-	return err
+	return finalName, err
 }
 
 // ProvisionOrClaimVizier provisions a given cluster or returns the ID if it already exists,
-func (s *Server) ProvisionOrClaimVizier(ctx context.Context, orgID uuid.UUID, userID uuid.UUID, clusterUID string, clusterName string) (uuid.UUID, error) {
+func (s *Server) ProvisionOrClaimVizier(ctx context.Context, orgID uuid.UUID, userID uuid.UUID, clusterUID string, clusterName string) (uuid.UUID, string, error) {
 	// TODO(zasgar): This duplicates some functionality in the Create function. Will deprecate that Create function soon.
 	tx, err := s.db.BeginTxx(ctx, nil)
 	if err != nil {
-		return uuid.Nil, vzerrors.ErrInternalDB
+		return uuid.Nil, "", vzerrors.ErrInternalDB
 	}
 	defer tx.Rollback()
 
@@ -1007,20 +1007,20 @@ func (s *Server) ProvisionOrClaimVizier(ctx context.Context, orgID uuid.UUID, us
 		return name
 	}
 
-	assignNameAndCommit := func() (uuid.UUID, error) {
+	assignNameAndCommit := func() (uuid.UUID, string, error) {
 		// Check if cluster already has a name.
 		var existingName *string
 
 		query := `SELECT cluster_name from vizier_cluster WHERE id=$1`
 		err := tx.QueryRowxContext(ctx, query, clusterID).Scan(&existingName)
 		if err != nil {
-			return uuid.Nil, vzerrors.ErrInternalDB
+			return uuid.Nil, "", vzerrors.ErrInternalDB
 		}
 
 		if existingName != nil {
 			// No input name specified, so no need to change cluster name.
 			if inputName == "" {
-				return clusterID, nil
+				return clusterID, *existingName, nil
 			}
 
 			// The existing name is already the same as the input name, or a derivation
@@ -1031,14 +1031,14 @@ func (s *Server) ProvisionOrClaimVizier(ctx context.Context, orgID uuid.UUID, us
 			// cannot distinguish between randomly generated names and actual-unaltered names.
 			dbName := *existingName
 			if inputName == dbName {
-				return clusterID, nil
+				return clusterID, *existingName, nil
 			}
 			prefixIndex := strings.LastIndex(dbName, "_")
 			if prefixIndex != -1 {
 				dbName = dbName[:prefixIndex]
 			}
 			if inputName == dbName {
-				return clusterID, nil
+				return clusterID, *existingName, nil
 			}
 		}
 
@@ -1047,13 +1047,14 @@ func (s *Server) ProvisionOrClaimVizier(ctx context.Context, orgID uuid.UUID, us
 			generateNameFunc = generateFromGivenName
 		}
 
-		if err := setClusterName(ctx, tx, clusterID, generateNameFunc); err != nil {
-			return uuid.Nil, vzerrors.ErrInternalDB
+		finalName, err := setClusterName(ctx, tx, clusterID, generateNameFunc)
+		if err != nil {
+			return uuid.Nil, "", vzerrors.ErrInternalDB
 		}
 
 		if err := tx.Commit(); err != nil {
 			log.WithError(err).Error("Failed to commit transaction")
-			return uuid.Nil, vzerrors.ErrInternalDB
+			return uuid.Nil, "", vzerrors.ErrInternalDB
 		}
 
 		events.Client().Enqueue(&analytics.Track{
@@ -1063,30 +1064,30 @@ func (s *Server) ProvisionOrClaimVizier(ctx context.Context, orgID uuid.UUID, us
 				Set("cluster_id", clusterID.String()).
 				Set("org_id", orgID.String()),
 		})
-		return clusterID, nil
+		return clusterID, finalName, nil
 	}
 
 	clusterID, status, err := findVizierWithUID(ctx, tx, orgID, clusterUID)
 	if err != nil {
-		return uuid.Nil, err
+		return uuid.Nil, "", err
 	}
 	if clusterID != uuid.Nil {
 		if status != vizierStatus(cvmsgspb.VZ_ST_DISCONNECTED) {
-			return uuid.Nil, vzerrors.ErrProvisionFailedVizierIsActive
+			return uuid.Nil, "", vzerrors.ErrProvisionFailedVizierIsActive
 		}
 		return assignNameAndCommit()
 	}
 
 	clusterID, _, err = findVizierWithEmptyUID(ctx, tx, orgID)
 	if err != nil {
-		return uuid.Nil, err
+		return uuid.Nil, "", err
 	}
 	if clusterID != uuid.Nil {
 		// Set the cluster ID.
 		query := `UPDATE vizier_cluster SET cluster_uid=$1 WHERE id=$2`
 		rows, err := tx.QueryxContext(ctx, query, clusterUID, clusterID)
 		if err != nil {
-			return uuid.Nil, err
+			return uuid.Nil, "", err
 		}
 		rows.Close()
 		return assignNameAndCommit()
@@ -1100,7 +1101,7 @@ func (s *Server) ProvisionOrClaimVizier(ctx context.Context, orgID uuid.UUID, us
 		INSERT INTO vizier_cluster_info(vizier_cluster_id, status) SELECT id, 'DISCONNECTED' FROM ins RETURNING vizier_cluster_id`
 	err = tx.QueryRowContext(ctx, query, orgID, DefaultProjectName, clusterUID).Scan(&clusterID)
 	if err != nil {
-		return uuid.Nil, err
+		return uuid.Nil, "", err
 	}
 	return assignNameAndCommit()
 }