Merge main in and some err handling changes

armory-io · Mar 21, 2023 · 3e87c51 · 3e87c51
2 parents 89058e3 + 9a8d3ad
commit 3e87c51
Showing 1 changed file with 30 additions and 14 deletions.
diff --git a/eks-updater.go b/eks-updater.go
@@ -35,18 +35,35 @@ func main() {
 	if len(*clusterName) == 0 {
 		log.Fatal("Invalid cluster name!  Must be set!")
 	}
-	if len(*nodegroupName) == 0 {
-		log.Fatal("Invalid nodegroup name!  Must be set!")
-	}
 
 	// Load the Shared AWS Configuration (~/.aws/config)
 	ctx := context.TODO()
-	client := getEksClient(ctx, *region, *roleArn)
+	client, err := getEksClient(ctx, *region, *roleArn)
+	if err != nil {
+		log.Fatal("Unable to get EKS client:", err)
+	}
+
 	log.Println("INFO: Starting updates...")
 	clusterInformation, _ := client.DescribeCluster(ctx, &eks.DescribeClusterInput{Name: clusterName})
-	updateError := updateClusterNodeGroup(client, ctx, clusterName, nodegroupName, waitTimeForNodeUpdates)
-	if updateError != nil {
-		log.Fatal("ERROR: Unable to update cluster node group... ", updateError)
+	if len(*nodegroupName) == 0 {
+		// Lookup and update the node groups...
+		nodeGroups, nodeGroupListErr := client.ListNodegroups(ctx, &eks.ListNodegroupsInput{ClusterName: clusterName})
+		if nodeGroupListErr != nil {
+			log.Fatal("ERROR: Unable to list node groups... ", nodeGroupListErr)
+		}
+		for _, nodeGroup := range nodeGroups.Nodegroups {
+			log.Println("INFO: Starting updates of node group " + nodeGroup)
+			updateError := updateClusterNodeGroup(client, ctx, clusterName, &nodeGroup, waitTimeForNodeUpdates)
+			if updateError != nil {
+				log.Fatal("ERROR: Unable to update cluster node group... ", updateError)
+			}
+		}
+	} else {
+		log.Println("INFO: Starting updates of node group " + *nodegroupName)
+		updateError := updateClusterNodeGroup(client, ctx, clusterName, nodegroupName, waitTimeForNodeUpdates)
+		if updateError != nil {
+			log.Fatal("ERROR: Unable to update cluster node group... ", updateError)
+		}
 	}
 	for _, addon := range addonsToUpdate {
 		updateAddon(client, ctx, clusterName, &addon, clusterInformation.Cluster.Version)
@@ -110,17 +127,16 @@ func updateClusterNodeGroup(client *eks.Client, ctx context.Context, clusterName
 	return nil
 }
 
-func getEksClient(ctx context.Context, region string, roleArn string) *eks.Client {
+func getEksClient(ctx context.Context, region string, roleArn string) (client *eks.Client, err error) {
 
 	var cfg aws.Config
-	var err error
 	cfg, err = config.LoadDefaultConfig(ctx, config.WithRegion(region))
 
 	if err != nil {
-		log.Fatal("ERROR: Unable to auth/get connected to AWS", err)
+		return client, err
 	}
 	if len(roleArn) == 0 {
-		return eks.NewFromConfig(cfg)
+		return eks.NewFromConfig(cfg), err
 	}
 	log.Println("INFO: Assuming role ARN " + roleArn)
 	// Create config & sts client with source account
@@ -136,14 +152,14 @@ func getEksClient(ctx context.Context, region string, roleArn string) *eks.Clien
 		DurationSeconds: &duration,
 	})
 	if err != nil {
-		log.Fatalf("unable to assume target role, %v", err)
+		return client, err
 	}
 	var assumedRoleCreds = response.Credentials
 
 	// Create config with target service client, using assumed role
 	cfg, err = config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(*assumedRoleCreds.AccessKeyId, *assumedRoleCreds.SecretAccessKey, *assumedRoleCreds.SessionToken)))
 	if err != nil {
-		log.Fatalf("unable to load static credentials for service client config, %v", err)
+		return client, err
 	}
-	return eks.NewFromConfig(cfg)
+	return eks.NewFromConfig(cfg), err
 }