support boskos resource of type image-builder-aws-account

cmd/aws-janitor-boskos/main.go

@@ -35,6 +35,7 @@ import (
 
 var (
 	boskosURL          = flag.String("boskos-url", "http://boskos", "Boskos URL")
+	rTypes             common.CommaSeparatedStrings
 	username           = flag.String("username", "", "Username used to access the Boskos server")
 	passwordFile       = flag.String("password-file", "", "The path to password file used to access the Boskos server")
 	region             = flag.String("region", "", "The region to clean (otherwise defaults to all regions)")
@@ -48,6 +49,10 @@ const (
 	sleepTime = time.Minute
 )
 
+func init() {
+	flag.Var(&rTypes, "resource-type", "comma-separated list of resources need to be cleaned up")
+}
+
 func main() {
 	logrusutil.ComponentInit()
 	flag.Parse()
@@ -64,6 +69,11 @@ func main() {
 		sweepSleepDuration = d
 	}
 
+	if len(rTypes) == 0 {
+		logrus.Info("--resource-type is empty! Setting it to default: aws-account")
+		rTypes = []string{"aws-account"}
+	}
+
 	boskos, err := client.NewClient("AWSJanitor", *boskosURL, *username, *passwordFile)
 	if err != nil {
 		logrus.WithError(err).Fatal("unable to create a Boskos client")
@@ -75,21 +85,23 @@ func main() {
 
 func run(boskos *client.Client) error {
 	for {
-		if res, err := boskos.Acquire(awsboskos.ResourceType, common.Dirty, common.Cleaning); errors.Cause(err) == client.ErrNotFound {
-			logrus.Info("no resource acquired. Sleeping.")
-			time.Sleep(sleepTime)
-			continue
-		} else if err != nil {
-			return errors.Wrap(err, "Couldn't retrieve resources from Boskos")
-		} else {
-			logrus.WithField("name", res.Name).Info("Acquired resource")
-			if err := cleanResource(res); err != nil {
-				return errors.Wrapf(err, "Couldn't clean resource %q", res.Name)
-			}
-			if err := boskos.ReleaseOne(res.Name, common.Free); err != nil {
-				return errors.Wrapf(err, "Failed to release resoures %q", res.Name)
+		for _, resourceType := range rTypes {
+			if res, err := boskos.Acquire(resourceType, common.Dirty, common.Cleaning); errors.Cause(err) == client.ErrNotFound {
+				logrus.Info("no resource acquired. Sleeping.")
+				time.Sleep(sleepTime)
+				continue
+			} else if err != nil {
+				return errors.Wrap(err, "Couldn't retrieve resources from Boskos")
+			} else {
+				logrus.WithField("name", res.Name).Info("Acquired resource")
+				if err := cleanResource(res); err != nil {
+					return errors.Wrapf(err, "Couldn't clean resource %q", res.Name)
+				}
+				if err := boskos.ReleaseOne(res.Name, common.Free); err != nil {
+					return errors.Wrapf(err, "Failed to release resoures %q", res.Name)
+				}
+				logrus.WithField("name", res.Name).Info("Released resource")
 			}
-			logrus.WithField("name", res.Name).Info("Released resource")
 		}
 	}
 }

common/aws/common.go

@@ -19,15 +19,13 @@ package aws
 import (
 	"errors"
 	"fmt"
+	"strings"
 
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"sigs.k8s.io/boskos/common"
 )
 
 const (
-	// ResourceType is the Boskos Type for AWS accounts
-	ResourceType = "aws-account"
-
 	// UserDataAccessIDKey is the key in UserData containing the AWS Access key ID
 	UserDataAccessIDKey = "access-key-id"
 
@@ -39,8 +37,8 @@ const (
 func GetAWSCreds(r *common.Resource) (credentials.Value, error) {
 	val := credentials.Value{}
 
-	if r.Type != ResourceType {
-		return val, fmt.Errorf("Wanted resource of type %q, got %q", ResourceType, r.Type)
+	if !strings.HasSuffix(r.Type, "aws-account") {
+		return val, fmt.Errorf("invalid aws resource type %q", r.Type)
 	}
 
 	accessKey, ok := r.UserData.Map.Load(UserDataAccessIDKey)
@@ -56,5 +54,4 @@ func GetAWSCreds(r *common.Resource) (credentials.Value, error) {
 	val.SecretAccessKey = secretKey.(string)
 
 	return val, nil
-
 }