Clean up.

artie-labs · Aug 4, 2024 · 7759d05 · 7759d05
1 parent 5509dd9
commit 7759d05
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/main.go b/main.go
@@ -34,13 +34,13 @@ func setUpMetrics(cfg *config.Metrics) (mtr.Client, error) {
 	return client, nil
 }
 
-func buildSource(cfg *config.Settings) (sources.Source, bool, error) {
+func buildSource(ctx context.Context, cfg *config.Settings) (sources.Source, bool, error) {
 	var source sources.Source
 	var isStreamingMode bool
 	var err error
 	switch cfg.Source {
 	case config.SourceDynamo:
-		source, isStreamingMode, err = dynamodb.Load(*cfg.DynamoDB)
+		source, isStreamingMode, err = dynamodb.Load(ctx, *cfg.DynamoDB)
 	case config.SourceMongoDB:
 		return mongo.Load(*cfg.MongoDB)
 	case config.SourceMSSQL:
@@ -108,7 +108,7 @@ func main() {
 		logger.Fatal(fmt.Sprintf("Failed to init %q destination writer", cfg.Destination), slog.Any("err", err))
 	}
 
-	source, isStreamingMode, err := buildSource(cfg)
+	source, isStreamingMode, err := buildSource(ctx, cfg)
 	if err != nil {
 		logger.Fatal(fmt.Sprintf("Failed to init %q source", cfg.Source), slog.Any("err", err))
 	}

diff --git a/sources/dynamodb/dynamodb.go b/sources/dynamodb/dynamodb.go
@@ -28,7 +28,7 @@ const (
 	concurrencyLimit = 30
 )
 
-func Load(cfg config.DynamoDB) (sources.Source, bool, error) {
+func Load(ctx context.Context, cfg config.DynamoDB) (sources.Source, bool, error) {
 	// TODO: Use v2
 	sess, err := session.NewSession(&aws.Config{
 		Region:      ptr.ToString(cfg.AwsRegion),
@@ -39,7 +39,8 @@ func Load(cfg config.DynamoDB) (sources.Source, bool, error) {
 	}
 
 	if cfg.Snapshot {
-		_awsCfg, err := awsCfg.LoadDefaultConfig(context.Background(),
+		_awsCfg, err := awsCfg.LoadDefaultConfig(
+			ctx,
 			awsCfg.WithRegion(cfg.AwsRegion),
 			awsCfg.WithCredentialsProvider(credentialsV2.NewStaticCredentialsProvider(cfg.AwsAccessKeyID, cfg.AwsSecretAccessKey, "")),
 		)