From 3cd4bdb2af8a1a6ef3f9efd7318294d453ec3b47 Mon Sep 17 00:00:00 2001
From: Ioannis Georgoulas <ioannis@paddle.com>
Date: Fri, 14 Aug 2020 13:29:50 +0100
Subject: [PATCH] Refactor get parameters by path call to not have 10 params
 limit

---
 parameter_store_client.go      | 22 +++++++++-------
 parameter_store_client_test.go | 48 ++++++++++++++++++++++++++++------
 2 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/parameter_store_client.go b/parameter_store_client.go
index 2a7b1a8..11c28b9 100644
--- a/parameter_store_client.go
+++ b/parameter_store_client.go
@@ -16,7 +16,7 @@ var (
 )
 
 type ssmClient interface {
-	GetParametersByPath(input *ssm.GetParametersByPathInput) (*ssm.GetParametersByPathOutput, error)
+	GetParametersByPathPages(input *ssm.GetParametersByPathInput, fn func(*ssm.GetParametersByPathOutput, bool) bool) error
 	GetParameter(input *ssm.GetParameterInput) (*ssm.GetParameterOutput, error)
 }
 
@@ -34,20 +34,22 @@ func (ps *ParameterStore) GetAllParametersByPath(path string, decrypt bool) (*Pa
 	var input = &ssm.GetParametersByPathInput{}
 	input.SetWithDecryption(decrypt)
 	input.SetPath(path)
+	input.SetMaxResults(10)
 	return ps.getParameters(input)
 }
 
 func (ps *ParameterStore) getParameters(input *ssm.GetParametersByPathInput) (*Parameters, error) {
-	result, err := ps.ssm.GetParametersByPath(input)
-	if err != nil {
-		return nil, err
-	}
-	parameters := NewParameters(*input.Path, make(map[string]*Parameter, len(result.Parameters)))
-	for _, v := range result.Parameters {
-		if v.Name == nil {
-			continue
+	parameters := NewParameters(*input.Path, make(map[string]*Parameter))
+	if err := ps.ssm.GetParametersByPathPages(input, func(result *ssm.GetParametersByPathOutput, b bool) bool {
+		for _, v := range result.Parameters {
+			if v.Name == nil {
+				continue
+			}
+			parameters.parameters[*v.Name] = &Parameter{Value: v.Value}
 		}
-		parameters.parameters[*v.Name] = &Parameter{Value: v.Value}
+		return !b
+	}); err != nil {
+		return nil, err
 	}
 	return parameters, nil
 }
diff --git a/parameter_store_client_test.go b/parameter_store_client_test.go
index aa6bf80..cf15eab 100644
--- a/parameter_store_client_test.go
+++ b/parameter_store_client_test.go
@@ -18,17 +18,35 @@ var param2 = new(ssm.Parameter).
 	SetValue("rds.something.aws.com").
 	SetARN("arn:aws:ssm:us-east-2:aws-account-id:/my-service/dev/DB_HOST")
 
+var param3 = new(ssm.Parameter).
+	SetName("/my-service/dev/DB_USERNAME").
+	SetValue("username").
+	SetARN("arn:aws:ssm:us-east-2:aws-account-id:/my-service/dev/DB_USERNAME")
+
 var errSSM = errors.New("ssm request error")
 
+type stubGetParametersByPathOutput struct {
+	Output         ssm.GetParametersByPathOutput
+	MoreParamsLeft bool
+}
+
 type stubSSMClient struct {
-	GetParametersByPathOutput *ssm.GetParametersByPathOutput
-	GetParametersByPathError  error
-	GetParameterOutput        *ssm.GetParameterOutput
-	GetParameterError         error
+	GetParametersByPathOutput        []stubGetParametersByPathOutput
+	GetParametersByPathError         error
+	GetParameterOutput               *ssm.GetParameterOutput
+	GetParameterError                error
 }
 
-func (s stubSSMClient) GetParametersByPath(input *ssm.GetParametersByPathInput) (*ssm.GetParametersByPathOutput, error) {
-	return s.GetParametersByPathOutput, s.GetParametersByPathError
+func (s stubSSMClient) GetParametersByPathPages(input *ssm.GetParametersByPathInput, fn func(*ssm.GetParametersByPathOutput, bool) bool) error {
+	if s.GetParametersByPathError == nil {
+		for _, output := range s.GetParametersByPathOutput {
+			done := fn(&output.Output, output.MoreParamsLeft)
+			if done {
+				return nil
+			}
+		}
+	}
+	return s.GetParametersByPathError
 }
 
 func (s stubSSMClient) GetParameter(input *ssm.GetParameterInput) (*ssm.GetParameterOutput, error) {
@@ -46,8 +64,21 @@ func TestClient_GetParametersByPath(t *testing.T) {
 		{
 			name: "Success",
 			ssmClient: &stubSSMClient{
-				GetParametersByPathOutput: &ssm.GetParametersByPathOutput{
-					Parameters: getParameters(),
+				GetParametersByPathOutput: []stubGetParametersByPathOutput{
+					{
+						MoreParamsLeft: true,
+						Output: ssm.GetParametersByPathOutput{
+							Parameters: getParameters(),
+						},
+					},
+					{
+						MoreParamsLeft: false,
+						Output: ssm.GetParametersByPathOutput{
+							Parameters: []*ssm.Parameter{
+								param3,
+							},
+						},
+					},
 				},
 			},
 			path: "/my-service/dev/",
@@ -56,6 +87,7 @@ func TestClient_GetParametersByPath(t *testing.T) {
 				parameters: map[string]*Parameter{
 					"/my-service/dev/DB_PASSWORD": {Value: param1.Value},
 					"/my-service/dev/DB_HOST":     {Value: param2.Value},
+					"/my-service/dev/DB_USERNAME": {Value: param3.Value},
 				},
 			},
 		},