weaver-localstack/src/main/scala/dev/pawlik/michal/aws.scala at master · majk-p/weaver-localstack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
package dev.pawlik.michal

import cats.effect.IO
import cats.implicits.*
import io.laserdisc.pure.kinesis.tagless.KinesisAsyncClientOp
import io.laserdisc.pure.s3.tagless.S3AsyncClientOp
import software.amazon.awssdk.core.async.AsyncRequestBody
import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest
import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest
import software.amazon.awssdk.services.kinesis.model.ListShardsRequest
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType
import software.amazon.awssdk.services.s3.model.PutObjectRequest

import scala.jdk.CollectionConverters.*
import java.time.Instant

object aws {
  object extensions {

    extension (kinesisClient: KinesisAsyncClientOp[IO]) {
      def readMessages(streamName: String): IO[List[String]] =
        for {
          shardsResponse <- kinesisClient.listShards(
            ListShardsRequest.builder().streamName(streamName).build()
          )
          shards = shardsResponse.shards().asScala.toList
          iteratorsResponses <- shards.parTraverse { shard =>
            kinesisClient.getShardIterator(
              GetShardIteratorRequest
                .builder()
                .streamName(streamName)
                .shardIteratorType(ShardIteratorType.TRIM_HORIZON)
                .shardId(shard.shardId())
                .build()
            )
          }
          iterators = iteratorsResponses.map(_.shardIterator())
          recordsResponses <- iterators.parTraverse { iterator =>
            kinesisClient.getRecords(
              GetRecordsRequest.builder().shardIterator(iterator).build()
            )
          }
          messages = recordsResponses
            .flatMap(_.records().asScala)
            .map(_.data().asUtf8String())
        } yield messages
    }

    extension (s3Client: S3AsyncClientOp[IO]) {
      def putMessage(bucketName: String)(message: String) =
        for {
          objectKey <- IO.randomUUID.map(_.toString())
          request = PutObjectRequest
            .builder()
            .bucket(bucketName)
            .key(objectKey)
            .build()
          body = AsyncRequestBody.fromString(message)
          _ <- s3Client.putObject(request, body)
        } yield ()
    }
  }
}