From 9a24f900750eb2672441fb5e3d2c5442b5b30cd4 Mon Sep 17 00:00:00 2001
From: Taichi Sasaki <tchsskk@gmail.com>
Date: Thu, 7 Sep 2023 20:08:03 +0900
Subject: [PATCH] Handle the split result of GetParametersByPath

---
 providers/parameterstore/parameterstore.go    | 36 +++++++++++--------
 .../parameterstore/parameterstore_test.go     | 34 +++++++++++++++---
 2 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/providers/parameterstore/parameterstore.go b/providers/parameterstore/parameterstore.go
index 1c419095..04df5f04 100644
--- a/providers/parameterstore/parameterstore.go
+++ b/providers/parameterstore/parameterstore.go
@@ -111,22 +111,30 @@ func (ps *ParameterStore[T]) Read() (map[string]interface{}, error) {
 			}
 		}
 	case ssm.GetParametersByPathInput:
-		output, err := ps.client.GetParametersByPath(context.TODO(), &input, ps.config.OptFns...)
-		if err != nil {
-			return nil, err
-		}
-		for _, p := range output.Parameters {
-			// If there's a transformation callback, run it.
-			if ps.config.Callback != nil {
-				name, value := ps.config.Callback(*p.Name, *p.Value)
-				// If the callback blanked the key, it should be omitted.
-				if name == "" {
-					break
+		var nextToken *string
+		for {
+			input.NextToken = nextToken
+			output, err := ps.client.GetParametersByPath(context.TODO(), &input, ps.config.OptFns...)
+			if err != nil {
+				return nil, err
+			}
+			for _, p := range output.Parameters {
+				// If there's a transformation callback, run it.
+				if ps.config.Callback != nil {
+					name, value := ps.config.Callback(*p.Name, *p.Value)
+					// If the callback blanked the key, it should be omitted.
+					if name == "" {
+						break
+					}
+					mp[name] = value
+				} else {
+					mp[*p.Name] = *p.Value
 				}
-				mp[name] = value
-			} else {
-				mp[*p.Name] = *p.Value
 			}
+			if output.NextToken == nil {
+				break
+			}
+			nextToken = output.NextToken
 		}
 	}
 	// Unflatten only when a delimiter is specified.
diff --git a/providers/parameterstore/parameterstore_test.go b/providers/parameterstore/parameterstore_test.go
index 30fd6f96..7611b6d2 100644
--- a/providers/parameterstore/parameterstore_test.go
+++ b/providers/parameterstore/parameterstore_test.go
@@ -22,9 +22,25 @@ func TestParameterStore(t *testing.T) {
 		config.WithAPIOptions([]func(*middleware.Stack) error{
 			// Mock the SDK response using the middleware.
 			func(stack *middleware.Stack) error {
+				type key struct{}
+				err := stack.Initialize.Add(
+					middleware.InitializeMiddlewareFunc(
+						"MockInitialize",
+						func(ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler) (out middleware.InitializeOutput, metadata middleware.Metadata, err error) {
+							switch v := in.Parameters.(type) {
+							case *ssm.GetParametersByPathInput:
+								ctx = middleware.WithStackValue(ctx, key{}, v.NextToken)
+							}
+							return next.HandleInitialize(ctx, in)
+						},
+					), middleware.Before,
+				)
+				if err != nil {
+					return err
+				}
 				return stack.Finalize.Add(
 					middleware.FinalizeMiddlewareFunc(
-						"Mock",
+						"MockFinalize",
 						func(ctx context.Context, input middleware.FinalizeInput, handler middleware.FinalizeHandler) (middleware.FinalizeOutput, middleware.Metadata, error) {
 							switch awsmiddleware.GetOperationName(ctx) {
 							case "GetParameter":
@@ -52,8 +68,10 @@ func TestParameterStore(t *testing.T) {
 									},
 								}, middleware.Metadata{}, nil
 							case "GetParametersByPath":
-								return middleware.FinalizeOutput{
-									Result: &ssm.GetParametersByPathOutput{
+								var output ssm.GetParametersByPathOutput
+								if middleware.GetStackValue(ctx, key{}) == (*string)(nil) {
+									output = ssm.GetParametersByPathOutput{
+										NextToken: aws.String("nextToken"),
 										Parameters: []types.Parameter{
 											{
 												Name:  aws.String("prefix.parent1"),
@@ -63,13 +81,19 @@ func TestParameterStore(t *testing.T) {
 												Name:  aws.String("prefix.parent2.child1"),
 												Value: aws.String("bob"),
 											},
+										},
+									}
+								} else {
+									output = ssm.GetParametersByPathOutput{
+										Parameters: []types.Parameter{
 											{
 												Name:  aws.String("prefix.parent2.child2.grandchild1"),
 												Value: aws.String("carol"),
 											},
 										},
-									},
-								}, middleware.Metadata{}, nil
+									}
+								}
+								return middleware.FinalizeOutput{Result: &output}, middleware.Metadata{}, nil
 							default:
 								return middleware.FinalizeOutput{}, middleware.Metadata{}, nil
 							}