Add Retention config option

Author: jcxplorer

cwlogger.go

@@ -28,6 +28,12 @@ type Config struct {
 	// handled during a PutLogEvents API call and caused a log events to be
 	// dropped.
 	ErrorReporter func(err error)
+
+	// An optional log group retention time in days. This value is only taken into
+	// account when creating a log group that does not yet exist. Set to 0
+	// (default) for no retention policy. Refer to the PutRetentionPolicy API
+	// documentation for valid values.
+	Retention int
 }
 
 // A Logger represents a single CloudWatch Logs log group.
@@ -40,6 +46,7 @@ type Logger struct {
 	wg            sync.WaitGroup
 	done          chan bool
 	errorReporter func(err error)
+	retention     int
 }
 
 // New creates a new Logger.
@@ -67,6 +74,7 @@ func New(config *Config) (*Logger, error) {
 		errorReporter: errorReporter,
 		name:          &config.LogGroupName,
 		svc:           config.Client,
+		retention:     config.Retention,
 		prefix:        randomHex(32),
 		batcher:       newBatcher(),
 		done:          make(chan bool),
@@ -135,6 +143,12 @@ func (lg *Logger) createIfNotExists() error {
 			}
 		}
 	}
+	if lg.retention != 0 {
+		_, err = lg.svc.PutRetentionPolicy(&cloudwatchlogs.PutRetentionPolicyInput{
+			LogGroupName:    lg.name,
+			RetentionInDays: aws.Int64(int64(lg.retention)),
+		})
+	}
 	return err
 }
 

cwlogger_test.go

@@ -56,10 +56,43 @@ func TestCreatesGroupAndStream(t *testing.T) {
 	assert.True(t, logStreamCreated)
 }
 
+func TestCreatesRetentionPolicy(t *testing.T) {
+	logGroupCreated := false
+	retentionPolicyCreated := false
+	config := &Config{
+		LogGroupName: "test",
+		Retention:    90,
+	}
+
+	newLoggerWithServer(config, func(w http.ResponseWriter, r *http.Request) {
+		if action(r) == "CreateLogGroup" {
+			logGroupCreated = true
+		}
+		if action(r) == "PutRetentionPolicy" {
+			if !logGroupCreated {
+				assert.Fail(t, "CreateLogGroup must be called before PutRetentionPolicy")
+				var data PutRetentionPolicy
+				parseBody(r, &data)
+				assert.Equal(t, "test", data.LogGroupName)
+				assert.Equal(t, 90, data.RetentionInDays)
+			}
+			retentionPolicyCreated = true
+		}
+	})
+
+	assert.True(t, logGroupCreated)
+	assert.True(t, retentionPolicyCreated)
+}
+
 func TestHandlesExistingGroup(t *testing.T) {
 	logStreamCreated := false
+	retentionPolicyCreated := false
+	config := &Config{
+		LogGroupName: "test",
+		Retention:    30,
+	}
 
-	newLoggerWithServer(defaultConfig, func(w http.ResponseWriter, r *http.Request) {
+	newLoggerWithServer(config, func(w http.ResponseWriter, r *http.Request) {
 		if action(r) == "CreateLogGroup" {
 			w.WriteHeader(http.StatusBadRequest)
 			w.Write([]byte(`
@@ -72,9 +105,13 @@ func TestHandlesExistingGroup(t *testing.T) {
 		if action(r) == "CreateLogStream" {
 			logStreamCreated = true
 		}
+		if action(r) == "PutRetentionPolicy" {
+			retentionPolicyCreated = true
+		}
 	})
 
 	assert.True(t, logStreamCreated)
+	assert.False(t, retentionPolicyCreated)
 }
 
 func TestSendsLogsToCloudWatchLogs(t *testing.T) {
@@ -415,6 +452,22 @@ func TestLogStreamCreationFails(t *testing.T) {
 	assert.Nil(t, logger)
 }
 
+func TestPutRetentionPolicyFails(t *testing.T) {
+	client := newClientWithServer(func(w http.ResponseWriter, r *http.Request) {
+		if action(r) == "PutRetentionPolicy" {
+			w.WriteHeader(http.StatusInternalServerError)
+			w.Write([]byte(`{"__type": "ServiceUnavailableException"}`))
+		}
+	})
+	logger, err := New(&Config{
+		Client:       client,
+		LogGroupName: "test",
+		Retention:    180,
+	})
+	assert.Error(t, err)
+	assert.Nil(t, logger)
+}
+
 func TestIgnoresBatchItCannotRetry(t *testing.T) {
 	var calls int
 
@@ -499,6 +552,11 @@ type PutLogEvents struct {
 	LogEvents     []*LogEvent `json:"logEvents"`
 }
 
+type PutRetentionPolicy struct {
+	LogGroupName    string `json:"logGroupName"`
+	RetentionInDays string `json:"retentionInDays"`
+}
+
 type LogEvent struct {
 	Timestamp int64  `json:"timestamp"`
 	Message   string `json:"message"`