projectlabeller.go

Documentation: edge-infra.dev/pkg/sds/lib/ci/dsds-sandboxes

     1  package dsdssandboxes
     2  
     3  import (
     4  	"context"
     5  
     6  	compute "google.golang.org/api/compute/v1"
     7  )
     8  
     9  type setLabelMethod func(string, *compute.InstancesSetLabelsRequest) error
    10  type listVMsmethod func() (*compute.InstanceList, error)
    11  type progressHookMethod func(string)
    12  
    13  type ProjectLabeller struct {
    14  	Name         string
    15  	setLabels    setLabelMethod
    16  	listVMs      listVMsmethod
    17  	progressHook progressHookMethod
    18  }
    19  
    20  func ProjectLabellerFromConf(conf BillingConf, zone string) (*ProjectLabeller, error) {
    21  	return ProjectLabellerFromName(conf.Metadata.Name, zone)
    22  }
    23  
    24  func ProjectLabellerFromName(name string, zone string) (*ProjectLabeller, error) {
    25  	ctx := context.Background()
    26  	service, err := compute.NewService(ctx)
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  
    31  	//Wrap calls to GCP in a closure to allow test injection
    32  	setLabels := func(instance string, labelRequest *compute.InstancesSetLabelsRequest) error {
    33  		_, err := service.Instances.SetLabels(name, zone, instance, labelRequest).Context(ctx).Do()
    34  		return err
    35  	}
    36  
    37  	listVMs := func() (*compute.InstanceList, error) {
    38  		instanceList, err := service.Instances.List(name, zone).Context(ctx).Do()
    39  		return instanceList, err
    40  	}
    41  
    42  	return &ProjectLabeller{
    43  		Name:         name,
    44  		setLabels:    setLabels,
    45  		listVMs:      listVMs,
    46  		progressHook: nil,
    47  	}, nil
    48  }
    49  
    50  // Provide a method that will be called with the name of the VM instance being processed
    51  func (p *ProjectLabeller) SetVMProgressHook(vmHook progressHookMethod) {
    52  	p.progressHook = vmHook
    53  }
    54  
    55  // Apply the schedule labels for the provided schedule to all VMs in this project
    56  func (p *ProjectLabeller) AddVMScheduleLabels(schedule Schedule) error {
    57  	instanceList, err := p.listVMs()
    58  	if err != nil {
    59  		return err
    60  	}
    61  
    62  	for _, instance := range instanceList.Items {
    63  		if p.progressHook != nil {
    64  			p.progressHook(instance.Name)
    65  		}
    66  
    67  		labels := schedule.ToLabelMap()
    68  		err := p.addLabelToVM(instance, labels)
    69  		if err != nil {
    70  			return err
    71  		}
    72  	}
    73  
    74  	return nil
    75  }
    76  
    77  func (p *ProjectLabeller) addLabelToVM(instance *compute.Instance, labels map[string]string) error {
    78  	fingerprint := instance.LabelFingerprint
    79  	combinedLabels := instance.Labels
    80  
    81  	for k, v := range labels {
    82  		combinedLabels[k] = v
    83  	}
    84  
    85  	labelRequest := &compute.InstancesSetLabelsRequest{
    86  		Labels:           combinedLabels,
    87  		LabelFingerprint: fingerprint,
    88  	}
    89  
    90  	err := p.setLabels(instance.Name, labelRequest)
    91  	if err != nil {
    92  		return err
    93  	}
    94  	return nil
    95  }
    96
View as plain text