github.com/pachyderm/pachyderm@v1.13.4/etc/testing/cleanup.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/json"
     5  	"os"
     6  	"os/exec"
     7  	"strings"
     8  	"time"
     9  
    10  	"github.com/aws/aws-lambda-go/lambda"
    11  	"github.com/aws/aws-sdk-go/aws"
    12  	"github.com/aws/aws-sdk-go/aws/session"
    13  	"github.com/aws/aws-sdk-go/service/s3"
    14  )
    15  
    16  // ClusterInfo holds information about a cluster.
    17  type ClusterInfo struct {
    18  	KopsBucket      string `json:"kops_bucket"`
    19  	PachydermBucket string `json:"pachyderm_bucket"`
    20  	Created         string `json:"created"`
    21  }
    22  
    23  // KopsBucket is the s3 bucket used by kops.
    24  const KopsBucket = "pachyderm-travis-state-store-v1"
    25  
    26  // MaxClusterTime is the maximimum time a cluster can be up.
    27  const MaxClusterTime = time.Hour * 4
    28  
    29  // HandleRequest handles the deletion of old clusters.
    30  func HandleRequest() (string, error) {
    31  	cmd := exec.Command("/bin/bash", "-c", "export PATH=$PATH:/var/task; kops --state=s3://"+KopsBucket+" get clusters | tail -n+2 | awk '{print $1}'")
    32  	cmd.Stderr = os.Stderr
    33  	out, err := cmd.Output()
    34  	if err != nil {
    35  		return "Failed to get clusters", err
    36  	}
    37  	names := strings.Split(string(out), "\n")
    38  	names = names[:len(names)-1]
    39  	var deleted string
    40  	s, err := session.NewSession()
    41  	if err != nil {
    42  		return "Failed to create session", err
    43  	}
    44  	svc := s3.New(s)
    45  	for _, name := range names {
    46  		infoObject, err := svc.GetObject(
    47  			&s3.GetObjectInput{
    48  				Bucket: aws.String(KopsBucket),
    49  				Key:    aws.String(name + "-info.json"),
    50  			})
    51  		if err != nil {
    52  			return "Failed to get info file", err
    53  		}
    54  		var info ClusterInfo
    55  		if err := json.NewDecoder(infoObject.Body).Decode(&info); err != nil {
    56  			return "Failed to decode info file", err
    57  		}
    58  		createTime, err := time.Parse(time.UnixDate, info.Created)
    59  		if err != nil {
    60  			return "Failed to parse create time", err
    61  		}
    62  		// Cluster has been up for too long
    63  		if createTime.Add(MaxClusterTime).Before(time.Now()) {
    64  			deleted += name + ", "
    65  			cmd := exec.Command("/bin/bash", "-c", "export PATH=$PATH:/var/task; kops --state=s3://"+KopsBucket+" delete cluster --name="+name+" --yes")
    66  			cmd.Stderr = os.Stderr
    67  			if err := cmd.Run(); err != nil {
    68  				return "Failed to delete cluster", err
    69  			}
    70  			_, err := svc.DeleteBucket(
    71  				&s3.DeleteBucketInput{
    72  					Bucket: aws.String(info.PachydermBucket),
    73  				})
    74  			if err != nil {
    75  				return "Failed to delete pachyderm bucket", err
    76  			}
    77  			_, err = svc.DeleteObject(
    78  				&s3.DeleteObjectInput{
    79  					Bucket: aws.String(KopsBucket),
    80  					Key:    aws.String(name + "-info.json"),
    81  				})
    82  			if err != nil {
    83  				return "Failed to delete info file", err
    84  			}
    85  		}
    86  	}
    87  	if len(deleted) <= 0 {
    88  		return "No clusters deleted", nil
    89  	}
    90  	deleted = "Clusters deleted: " + deleted
    91  	tmp := []rune(deleted)
    92  	return string(tmp[:len(tmp)-2]), nil
    93  }
    94  
    95  func main() {
    96  	lambda.Start(HandleRequest)
    97  }