wip

Author: mmetc

go.mod

@@ -14,7 +14,9 @@ require (
 	github.com/aws/aws-sdk-go v1.52.0
 	github.com/aws/aws-sdk-go-v2 v1.38.3
 	github.com/aws/aws-sdk-go-v2/config v1.31.6
+	github.com/aws/aws-sdk-go-v2/credentials v1.18.10 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.19.4
+	github.com/aws/aws-sdk-go-v2/service/kinesis v1.40.1
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.87.3
 	github.com/aws/aws-sdk-go-v2/service/sqs v1.42.3
 	github.com/beevik/etree v1.4.1
@@ -107,7 +109,6 @@ require (
 	github.com/apparentlymart/go-textseg/v15 v15.0.0 // indirect
 	github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.1 // indirect
-	github.com/aws/aws-sdk-go-v2/credentials v1.18.10 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.6 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.6 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.6 // indirect

go.sum

@@ -76,6 +76,8 @@ github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.6 h1:LHS1YAIJX
 github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.6/go.mod h1:c9PCiTEuh0wQID5/KqA32J+HAgZxN9tOGXKCiYJjTZI=
 github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.6 h1:nEXUSAwyUfLTgnc9cxlDWy637qsq4UWwp3sNAfl0Z3Y=
 github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.6/go.mod h1:HGzIULx4Ge3Do2V0FaiYKcyKzOqwrhUZgCI77NisswQ=
+github.com/aws/aws-sdk-go-v2/service/kinesis v1.40.1 h1:9QC0AF6gakV1TZuGp3NEUNl/6gXt3rfIifnxd+dWwbw=
+github.com/aws/aws-sdk-go-v2/service/kinesis v1.40.1/go.mod h1:UpSQbmXxFiDGDrvqsTgEm3YijDf9cg/Ti+s2W0SeFEU=
 github.com/aws/aws-sdk-go-v2/service/s3 v1.87.3 h1:ETkfWcXP2KNPLecaDa++5bsQhCRa5M5sLUJa5DWYIIg=
 github.com/aws/aws-sdk-go-v2/service/s3 v1.87.3/go.mod h1:+/3ZTqoYb3Ur7DObD00tarKMLMuKg8iqz5CHEanqTnw=
 github.com/aws/aws-sdk-go-v2/service/sqs v1.42.3 h1:0dWg1Tkz3FnEo48DgAh7CT22hYyMShly8WMd3sGx0xI=

pkg/acquisition/modules/kinesis/kinesis.go

@@ -11,10 +11,12 @@ import (
 	"strings"
 	"time"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/arn"
-	"github.com/aws/aws-sdk-go/aws/session"
-	"github.com/aws/aws-sdk-go/service/kinesis"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/aws/arn"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/kinesis"
+	kinTypes "github.com/aws/aws-sdk-go-v2/service/kinesis/types"
+
 	yaml "github.com/goccy/go-yaml"
 	"github.com/prometheus/client_golang/prometheus"
 	log "github.com/sirupsen/logrus"
@@ -45,7 +47,7 @@ type KinesisSource struct {
 	metricsLevel    metrics.AcquisitionMetricsLevel
 	Config          KinesisConfiguration
 	logger          *log.Entry
-	kClient         *kinesis.Kinesis
+	kClient         *kinesis.Client
 	shardReaderTomb *tomb.Tomb
 }
 
@@ -68,39 +70,32 @@ func (k *KinesisSource) GetUuid() string {
 	return k.Config.UniqueId
 }
 
-func (k *KinesisSource) newClient() error {
-	var sess *session.Session
-
-	if k.Config.AwsProfile != nil {
-		sess = session.Must(session.NewSessionWithOptions(session.Options{
-			SharedConfigState: session.SharedConfigEnable,
-			Profile:           *k.Config.AwsProfile,
-		}))
-	} else {
-		sess = session.Must(session.NewSessionWithOptions(session.Options{
-			SharedConfigState: session.SharedConfigEnable,
-		}))
+func (k *KinesisSource) newClient(ctx context.Context) error {
+	var loadOpts []func(*config.LoadOptions) error
+	if k.Config.AwsProfile != nil && *k.Config.AwsProfile != "" {
+		loadOpts = append(loadOpts, config.WithSharedConfigProfile(*k.Config.AwsProfile))
 	}
-
-	if sess == nil {
-		return errors.New("failed to create aws session")
+	region := k.Config.AwsRegion
+	if region == "" {
+		region = "us-east-1"
 	}
+	loadOpts = append(loadOpts, config.WithRegion(region))
+	loadOpts = append(loadOpts, config.WithCredentialsProvider(aws.AnonymousCredentials{}))
 
-	config := aws.NewConfig()
-
-	if k.Config.AwsRegion != "" {
-		config = config.WithRegion(k.Config.AwsRegion)
+	cfg, err := config.LoadDefaultConfig(ctx, loadOpts...)
+	if err != nil {
+		return fmt.Errorf("failed to load aws config: %w", err)
 	}
 
+	var clientOpts []func(*kinesis.Options)
 	if k.Config.AwsEndpoint != "" {
-		config = config.WithEndpoint(k.Config.AwsEndpoint)
+		clientOpts = append(clientOpts, func(o *kinesis.Options) {
+			o.BaseEndpoint = aws.String(k.Config.AwsEndpoint)
+		})
 	}
 
-	k.kClient = kinesis.New(sess, config)
-	if k.kClient == nil {
-		return errors.New("failed to create kinesis client")
-	}
 
+	k.kClient = kinesis.NewFromConfig(cfg, clientOpts...)
 	return nil
 }
 
@@ -156,7 +151,7 @@ func (k *KinesisSource) Configure(yamlConfig []byte, logger *log.Entry, metricsL
 		return err
 	}
 
-	err = k.newClient()
+	err = k.newClient(context.TODO())
 	if err != nil {
 		return fmt.Errorf("cannot create kinesis client: %w", err)
 	}
@@ -203,12 +198,14 @@ func (k *KinesisSource) decodeFromSubscription(record []byte) ([]CloudwatchSubsc
 func (k *KinesisSource) WaitForConsumerDeregistration(consumerName string, streamARN string) error {
 	maxTries := k.Config.MaxRetries
 	for i := range maxTries {
-		_, err := k.kClient.DescribeStreamConsumer(&kinesis.DescribeStreamConsumerInput{
-			ConsumerName: aws.String(consumerName),
-			StreamARN:    aws.String(streamARN),
-		})
+		_, err := k.kClient.DescribeStreamConsumer(
+			context.TODO(),
+			&kinesis.DescribeStreamConsumerInput{
+				ConsumerName: aws.String(consumerName),
+				StreamARN:    aws.String(streamARN),
+			})
 
-		var resourceNotFoundErr *kinesis.ResourceNotFoundException
+		var resourceNotFoundErr *kinTypes.ResourceNotFoundException
 		if errors.As(err, &resourceNotFoundErr) {
 			return nil
 		}
@@ -226,12 +223,14 @@ func (k *KinesisSource) WaitForConsumerDeregistration(consumerName string, strea
 
 func (k *KinesisSource) DeregisterConsumer() error {
 	k.logger.Debugf("Deregistering consumer %s if it exists", k.Config.ConsumerName)
-	_, err := k.kClient.DeregisterStreamConsumer(&kinesis.DeregisterStreamConsumerInput{
-		ConsumerName: aws.String(k.Config.ConsumerName),
-		StreamARN:    aws.String(k.Config.StreamARN),
-	})
+	_, err := k.kClient.DeregisterStreamConsumer(
+		context.TODO(),
+		&kinesis.DeregisterStreamConsumerInput{
+			ConsumerName: aws.String(k.Config.ConsumerName),
+			StreamARN:    aws.String(k.Config.StreamARN),
+		})
 
-	var resourceNotFoundErr *kinesis.ResourceNotFoundException
+	var resourceNotFoundErr *kinTypes.ResourceNotFoundException
 	if errors.As(err, &resourceNotFoundErr) {
 		return nil
 	}
@@ -251,14 +250,16 @@ func (k *KinesisSource) DeregisterConsumer() error {
 func (k *KinesisSource) WaitForConsumerRegistration(consumerARN string) error {
 	maxTries := k.Config.MaxRetries
 	for i := range maxTries {
-		describeOutput, err := k.kClient.DescribeStreamConsumer(&kinesis.DescribeStreamConsumerInput{
-			ConsumerARN: aws.String(consumerARN),
-		})
+		describeOutput, err := k.kClient.DescribeStreamConsumer(
+			context.TODO(),
+				&kinesis.DescribeStreamConsumerInput{
+				ConsumerARN: aws.String(consumerARN),
+			})
 		if err != nil {
 			return fmt.Errorf("cannot describe stream consumer: %w", err)
 		}
 
-		if *describeOutput.ConsumerDescription.ConsumerStatus == "ACTIVE" {
+		if describeOutput.ConsumerDescription.ConsumerStatus == "ACTIVE" {
 			k.logger.Debugf("Consumer %s is active", consumerARN)
 			return nil
 		}
@@ -273,10 +274,12 @@ func (k *KinesisSource) WaitForConsumerRegistration(consumerARN string) error {
 func (k *KinesisSource) RegisterConsumer() (*kinesis.RegisterStreamConsumerOutput, error) {
 	k.logger.Debugf("Registering consumer %s", k.Config.ConsumerName)
 
-	streamConsumer, err := k.kClient.RegisterStreamConsumer(&kinesis.RegisterStreamConsumerInput{
-		ConsumerName: aws.String(k.Config.ConsumerName),
-		StreamARN:    aws.String(k.Config.StreamARN),
-	})
+	streamConsumer, err := k.kClient.RegisterStreamConsumer(
+		context.TODO(),
+		&kinesis.RegisterStreamConsumerInput{
+			ConsumerName: aws.String(k.Config.ConsumerName),
+			StreamARN:    aws.String(k.Config.StreamARN),
+		})
 	if err != nil {
 		return nil, fmt.Errorf("cannot register stream consumer: %w", err)
 	}
@@ -289,7 +292,7 @@ func (k *KinesisSource) RegisterConsumer() (*kinesis.RegisterStreamConsumerOutpu
 	return streamConsumer, nil
 }
 
-func (k *KinesisSource) ParseAndPushRecords(records []*kinesis.Record, out chan types.Event, logger *log.Entry, shardID string) {
+func (k *KinesisSource) ParseAndPushRecords(records []kinTypes.Record, out chan types.Event, logger *log.Entry, shardID string) {
 	for _, record := range records {
 		if k.Config.StreamARN != "" {
 			if k.metricsLevel != metrics.AcquisitionMetricsLevelNone {
@@ -365,38 +368,42 @@ func (k *KinesisSource) ReadFromSubscription(reader kinesis.SubscribeToShardEven
 				return nil
 			}
 
-			switch event := event.(type) {
-			case *kinesis.SubscribeToShardEvent:
-				k.ParseAndPushRecords(event.Records, out, logger, shardID)
-			case *kinesis.SubscribeToShardEventStreamUnknownEvent:
-				logger.Infof("got an unknown event, what to do ?")
+			switch et := event.(type) {
+			case *kinTypes.SubscribeToShardEventStreamMemberSubscribeToShardEvent:
+				k.ParseAndPushRecords(et.Value.Records, out, logger, shardID)
+			default:
+				logger.Infof("unhandled SubscribeToShard event: %T", et)
 			}
 		}
 	}
 }
 
 func (k *KinesisSource) SubscribeToShards(arn arn.ARN, streamConsumer *kinesis.RegisterStreamConsumerOutput, out chan types.Event) error {
-	shards, err := k.kClient.ListShards(&kinesis.ListShardsInput{
-		StreamName: aws.String(arn.Resource[7:]),
-	})
+	shards, err := k.kClient.ListShards(
+		context.TODO(),
+		&kinesis.ListShardsInput{
+			StreamName: aws.String(arn.Resource[7:]),
+		})
 	if err != nil {
 		return fmt.Errorf("cannot list shards for enhanced_read: %w", err)
 	}
 
 	for _, shard := range shards.Shards {
 		shardID := *shard.ShardId
 
-		r, err := k.kClient.SubscribeToShard(&kinesis.SubscribeToShardInput{
-			ShardId:          aws.String(shardID),
-			StartingPosition: &kinesis.StartingPosition{Type: aws.String(kinesis.ShardIteratorTypeLatest)},
-			ConsumerARN:      streamConsumer.Consumer.ConsumerARN,
-		})
+		r, err := k.kClient.SubscribeToShard(
+			context.TODO(),
+			&kinesis.SubscribeToShardInput{
+				ShardId:          aws.String(shardID),
+				StartingPosition: &kinTypes.StartingPosition{Type: kinTypes.ShardIteratorTypeLatest},
+				ConsumerARN:      streamConsumer.Consumer.ConsumerARN,
+			})
 		if err != nil {
 			return fmt.Errorf("cannot subscribe to shard: %w", err)
 		}
 
 		k.shardReaderTomb.Go(func() error {
-			return k.ReadFromSubscription(r.GetEventStream().Reader, out, shardID, arn.Resource[7:])
+			return k.ReadFromSubscription(r.GetStream().Reader, out, shardID, arn.Resource[7:])
 		})
 	}
 
@@ -463,11 +470,13 @@ func (k *KinesisSource) ReadFromShard(out chan types.Event, shardID string) erro
 	logger := k.logger.WithField("shard", shardID)
 	logger.Debugf("Starting to read shard")
 
-	sharIt, err := k.kClient.GetShardIterator(&kinesis.GetShardIteratorInput{
-		ShardId:           aws.String(shardID),
-		StreamName:        &k.Config.StreamName,
-		ShardIteratorType: aws.String(kinesis.ShardIteratorTypeLatest),
-	})
+	sharIt, err := k.kClient.GetShardIterator(
+		context.TODO(),
+		&kinesis.GetShardIteratorInput{
+			ShardId:           aws.String(shardID),
+			StreamName:        &k.Config.StreamName,
+			ShardIteratorType: kinTypes.ShardIteratorTypeLatest,
+		})
 	if err != nil {
 		logger.Errorf("Cannot get shard iterator: %s", err)
 		return fmt.Errorf("cannot get shard iterator: %w", err)
@@ -480,16 +489,17 @@ func (k *KinesisSource) ReadFromShard(out chan types.Event, shardID string) erro
 	for {
 		select {
 		case <-ticker.C:
-			records, err := k.kClient.GetRecords(&kinesis.GetRecordsInput{ShardIterator: it})
+			records, err := k.kClient.GetRecords(context.TODO(), &kinesis.GetRecordsInput{ShardIterator: it})
+			it = records.NextShardIterator
 
-			var throughputErr *kinesis.ProvisionedThroughputExceededException
+			var throughputErr *kinTypes.ProvisionedThroughputExceededException
 			if errors.As(err, &throughputErr) {
 				logger.Warn("Provisioned throughput exceeded")
 				// TODO: implement exponential backoff
 				continue
 			}
 
-			var expiredIteratorErr *kinesis.ExpiredIteratorException
+			var expiredIteratorErr *kinTypes.ExpiredIteratorException
 			if errors.As(err, &expiredIteratorErr) {
 				logger.Warn("Expired iterator")
 				continue
@@ -521,9 +531,11 @@ func (k *KinesisSource) ReadFromStream(out chan types.Event, t *tomb.Tomb) error
 	k.logger.Info("starting kinesis acquisition from shards")
 
 	for {
-		shards, err := k.kClient.ListShards(&kinesis.ListShardsInput{
-			StreamName: aws.String(k.Config.StreamName),
-		})
+		shards, err := k.kClient.ListShards(
+			context.TODO(),
+			&kinesis.ListShardsInput{
+				StreamName: aws.String(k.Config.StreamName),
+			})
 		if err != nil {
 			return fmt.Errorf("cannot list shards: %w", err)
 		}

pkg/acquisition/modules/kinesis/kinesis_test.go

@@ -10,9 +10,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/session"
-	"github.com/aws/aws-sdk-go/service/kinesis"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/kinesis"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -53,8 +54,14 @@ func GenSubObject(t *testing.T, i int) []byte {
 }
 
 func WriteToStream(t *testing.T, endpoint string, streamName string, count int, shards int, sub bool) {
-	sess := session.Must(session.NewSession())
-	kinesisClient := kinesis.New(sess, aws.NewConfig().WithEndpoint(endpoint).WithRegion("us-east-1"))
+	ctx := t.Context()
+
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion("us-east-1"), config.WithCredentialsProvider(aws.AnonymousCredentials{}))
+	require.NoError(t, err)
+
+	kinesisClient := kinesis.NewFromConfig(cfg, func(o *kinesis.Options) {
+		o.BaseEndpoint = aws.String(endpoint)
+	})
 
 	for i := range count {
 		partition := "partition"
@@ -70,7 +77,7 @@ func WriteToStream(t *testing.T, endpoint string, streamName string, count int,
 			data = []byte(strconv.Itoa(i))
 		}
 
-		_, err := kinesisClient.PutRecord(&kinesis.PutRecordInput{
+		_, err := kinesisClient.PutRecord(ctx, &kinesis.PutRecordInput{
 			Data:         data,
 			PartitionKey: aws.String(partition),
 			StreamName:   aws.String(streamName),

pkg/acquisition/modules/s3/s3.go

@@ -143,7 +143,9 @@ func (s *S3Source) newS3Client() error {
 
 	var clientOpts []func(*s3.Options)
 	if s.Config.AwsEndpoint != "" {
-		clientOpts = append(clientOpts, func(o *s3.Options) { o.BaseEndpoint = aws.String(s.Config.AwsEndpoint) })
+		clientOpts = append(clientOpts, func(o *s3.Options) {
+			o.BaseEndpoint = aws.String(s.Config.AwsEndpoint)
+		})
 	}
 
 	s.s3Client = s3.NewFromConfig(cfg, clientOpts...)