From af3384eee412d26c99873d560043081d6e218b14 Mon Sep 17 00:00:00 2001
From: rho song <13165182+rhoboat@users.noreply.github.com>
Date: Wed, 28 Sep 2022 11:44:11 -0700
Subject: [PATCH] Allow logger injection.

---
 awscommons/v2/auth.go |  3 +++
 awscommons/v2/ecs.go  | 17 ++++++++---------
 logging/logging.go    |  6 ------
 3 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/awscommons/v2/auth.go b/awscommons/v2/auth.go
index a6d2ca8..3195525 100644
--- a/awscommons/v2/auth.go
+++ b/awscommons/v2/auth.go
@@ -6,6 +6,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/gruntwork-io/go-commons/errors"
+	"github.com/sirupsen/logrus"
 )
 
 const (
@@ -17,6 +18,8 @@ type Options struct {
 	Region string
 
 	Context context.Context
+
+	Logger logrus.Entry
 }
 
 // NewOptions will create a new aws.Options struct that provides reasonable defaults for unspecified values.
diff --git a/awscommons/v2/ecs.go b/awscommons/v2/ecs.go
index a45837d..b5a7d67 100644
--- a/awscommons/v2/ecs.go
+++ b/awscommons/v2/ecs.go
@@ -11,7 +11,6 @@ import (
 	ecsTypes "github.com/aws/aws-sdk-go-v2/service/ecs/types"
 	"github.com/gruntwork-io/go-commons/collections"
 	"github.com/gruntwork-io/go-commons/errors"
-	"github.com/gruntwork-io/go-commons/logging"
 	"github.com/gruntwork-io/go-commons/retry"
 )
 
@@ -29,7 +28,7 @@ func GetContainerInstanceArns(opts *Options, clusterName string) ([]string, erro
 		return nil, err
 	}
 
-	logger := logging.GetProjectLogger()
+	logger := opts.Logger
 	logger.Infof("Looking up Container Instance ARNs for ECS cluster %s", clusterName)
 
 	input := &ecs.ListContainerInstancesInput{Cluster: aws.String(clusterName)}
@@ -60,7 +59,7 @@ func StartDrainingContainerInstances(opts *Options, clusterName string, containe
 		return err
 	}
 
-	logger := logging.GetProjectLogger()
+	logger := opts.Logger
 	batchSize := 10
 	numBatches := int(math.Ceil(float64(len(containerInstanceArns) / batchSize)))
 
@@ -100,7 +99,7 @@ func WaitForContainerInstancesToDrain(opts *Options, clusterName string, contain
 		return err
 	}
 
-	logger := logging.GetProjectLogger()
+	logger := opts.Logger
 	logger.Infof("Checking if all ECS Tasks have been drained from the ECS Container Instances in Cluster %s.", clusterName)
 
 	batchSize := 100
@@ -133,7 +132,7 @@ func WaitForContainerInstancesToDrain(opts *Options, clusterName string, contain
 			}
 
 			// Yay, all done.
-			if drained, _ := allInstancesFullyDrained(responses); drained == true {
+			if drained, _ := allInstancesFullyDrained(opts, responses); drained == true {
 				logger.Infof("All container instances have been drained in Cluster %s!", clusterName)
 				return nil
 			}
@@ -164,7 +163,7 @@ func NewECSClient(opts *Options) (*ecs.Client, error) {
 	return ecs.NewFromConfig(cfg), nil
 }
 
-func allInstancesFullyDrained(responses []*ecs.DescribeContainerInstancesOutput) (bool, error) {
+func allInstancesFullyDrained(opts *Options, responses []*ecs.DescribeContainerInstancesOutput) (bool, error) {
 	for _, response := range responses {
 		instances := response.ContainerInstances
 		if len(instances) == 0 {
@@ -172,7 +171,7 @@ func allInstancesFullyDrained(responses []*ecs.DescribeContainerInstancesOutput)
 		}
 
 		for _, instance := range instances {
-			if !instanceFullyDrained(instance) {
+			if !instanceFullyDrained(opts, instance) {
 				return false, nil
 			}
 		}
@@ -180,8 +179,8 @@ func allInstancesFullyDrained(responses []*ecs.DescribeContainerInstancesOutput)
 	return true, nil
 }
 
-func instanceFullyDrained(instance ecsTypes.ContainerInstance) bool {
-	logger := logging.GetProjectLogger()
+func instanceFullyDrained(opts *Options, instance ecsTypes.ContainerInstance) bool {
+	logger := opts.Logger
 	instanceArn := instance.ContainerInstanceArn
 
 	if *instance.Status == "ACTIVE" {
diff --git a/logging/logging.go b/logging/logging.go
index 8ccb90e..d29d764 100644
--- a/logging/logging.go
+++ b/logging/logging.go
@@ -9,12 +9,6 @@ import (
 var globalLogLevel = logrus.InfoLevel
 var globalLogLevelLock = sync.Mutex{}
 
-// GetProjectLogger creates a new project logger
-func GetProjectLogger() *logrus.Entry {
-	logger := GetLogger("")
-	return logger.WithField("name", "go-commons")
-}
-
 // GetLogger create a new logger with the given name
 func GetLogger(name string) *logrus.Logger {
 	logger := logrus.New()