Added pagination support (#920)

zio · Aug 18, 2022 · 2b76747 · 2b76747
1 parent 5809fe0
commit 2b76747
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 7 deletions.
diff --git a/aws/shared/src/main/scala/zio/config/aws/parameterstore/ParameterStoreConfigSource.scala b/aws/shared/src/main/scala/zio/config/aws/parameterstore/ParameterStoreConfigSource.scala
@@ -6,7 +6,8 @@ import com.amazonaws.services.simplesystemsmanagement.{
   AWSSimpleSystemsManagementClientBuilder
 }
 import zio.config.{PropertyTreePath, ReadError, _}
-import zio.{Task, ZIO, ZManaged}
+import zio.stream.ZStream
+import zio.{Chunk, Task, ZIO, ZManaged}
 
 import scala.jdk.CollectionConverters._
 
@@ -28,12 +29,34 @@ object ParameterStoreConfigSource {
                   .withRecursive(true)
                   .withWithDecryption(true)
 
-              ZIO
-                .effect(ssm.getParametersByPath(request).getParameters)
-                .map(_.asScala.toList)
-                .map { list =>
+              ZStream
+                .paginateM(
+                  ZIO.effect(ssm.getParametersByPath(request))
+                )(_.map { response =>
+                  val currentBatchResult =
+                    Chunk.fromIterable(
+                      response.getParameters.asScala.toList
+                    )
+                  val nextToken          = response.getNextToken
+                  val nextBatch          =
+                    if (nextToken == null || nextToken.trim.isEmpty)
+                      None
+                    else
+                      Some(
+                        ZIO.effect(
+                          ssm.getParametersByPath(request.withNextToken(nextToken))
+                        )
+                      )
+
+                  (currentBatchResult, nextBatch)
+                })
+                .runCollect
+                .map { result =>
                   ConfigSource
-                    .getPropertyTreeFromMap(convertParameterListToMap(list, basePath), keyDelimiter = Some('/'))
+                    .getPropertyTreeFromMap(
+                      convertParameterListToMap(result.flatten.toList, basePath),
+                      keyDelimiter = Some('/')
+                    )
                 }
             }
               .map(tree => (path: PropertyTreePath[String]) => ZIO.succeed(tree.at(path)))

diff --git a/build.sbt b/build.sbt
@@ -205,6 +205,7 @@ lazy val zioConfigAws    = crossProject(JVMPlatform)
   .settings(
     libraryDependencies ++= Seq(
       "com.amazonaws" % "aws-java-sdk-ssm" % awsVersion,
+      "dev.zio"      %% "zio-streams"      % zioVersion,
       "dev.zio"      %% "zio-test"         % zioVersion % Test,
       "dev.zio"      %% "zio-test-sbt"     % zioVersion % Test
     ),

diff --git a/project/build.properties b/project/build.properties
@@ -1 +1 @@
-sbt.version = 1.5.7
+sbt.version = 1.7.1