Pull context out of dynamo.KafkaMessage (#33)

lib/dynamo/message.go

@@ -1,15 +1,16 @@
 package dynamo
 
 import (
-	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/artie-labs/reader/config"
+	"strconv"
+	"time"
+
 	"github.com/artie-labs/transfer/lib/cdc/util"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/segmentio/kafka-go"
-	"strconv"
-	"time"
+
+	"github.com/artie-labs/reader/config"
 )
 
 type Message struct {
@@ -99,16 +100,11 @@ func (m *Message) artieMessage() (util.SchemaEventPayload, error) {
 	}, nil
 }
 
-func (m *Message) TopicName(ctx context.Context) (string, error) {
-	cfg := config.FromContext(ctx)
-	if cfg.Kafka == nil {
-		return "", fmt.Errorf("kafka config is nil")
-	}
-
-	return fmt.Sprintf("%s.%s", cfg.Kafka.TopicPrefix, m.tableName), nil
+func (m *Message) TopicName(kafkaCfg config.Kafka) string {
+	return fmt.Sprintf("%s.%s", kafkaCfg.TopicPrefix, m.tableName)
 }
 
-func (m *Message) KafkaMessage(ctx context.Context) (kafka.Message, error) {
+func (m *Message) KafkaMessage(kafkaCfg config.Kafka) (kafka.Message, error) {
 	msg, err := m.artieMessage()
 	if err != nil {
 		return kafka.Message{}, fmt.Errorf("failed to generate artie message, err: %v", err)
@@ -124,10 +120,7 @@ func (m *Message) KafkaMessage(ctx context.Context) (kafka.Message, error) {
 		return kafka.Message{}, err
 	}
 
-	topic, err := m.TopicName(ctx)
-	if err != nil {
-		return kafka.Message{}, err
-	}
+	topic := m.TopicName(kafkaCfg)
 
 	return kafka.Message{
 		Topic: topic,

sources/dynamodb/dynamodb.go

@@ -5,16 +5,17 @@ import (
 	"log/slog"
 	"time"
 
-	"github.com/artie-labs/reader/config"
-	"github.com/artie-labs/reader/lib/logger"
-	"github.com/artie-labs/reader/lib/s3lib"
-	"github.com/artie-labs/reader/sources/dynamodb/offsets"
 	"github.com/artie-labs/transfer/lib/ptr"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 	"github.com/aws/aws-sdk-go/service/dynamodbstreams"
+
+	"github.com/artie-labs/reader/config"
+	"github.com/artie-labs/reader/lib/logger"
+	"github.com/artie-labs/reader/lib/s3lib"
+	"github.com/artie-labs/reader/sources/dynamodb/offsets"
 )
 
 type Store struct {
@@ -84,7 +85,7 @@ func (s *Store) Run(ctx context.Context) {
 		go s.ListenToChannel(ctx)
 
 		// Scan it for the first time manually, so we don't have to wait 5 mins
-		s.scanForNewShards(ctx)
+		s.scanForNewShards()
 		for {
 			select {
 			case <-ctx.Done():
@@ -93,13 +94,13 @@ func (s *Store) Run(ctx context.Context) {
 				return
 			case <-ticker.C:
 				slog.Info("Scanning for new shards...")
-				s.scanForNewShards(ctx)
+				s.scanForNewShards()
 			}
 		}
 	}
 }
 
-func (s *Store) scanForNewShards(ctx context.Context) {
+func (s *Store) scanForNewShards() {
 	var exclusiveStartShardId *string
 	for {
 		input := &dynamodbstreams.DescribeStreamInput{

sources/dynamodb/shard.go

@@ -5,6 +5,7 @@ import (
 	"log/slog"
 	"time"
 
+	"github.com/artie-labs/reader/config"
 	"github.com/artie-labs/reader/lib/dynamo"
 	"github.com/artie-labs/reader/lib/kafkalib"
 	"github.com/artie-labs/reader/lib/logger"
@@ -64,6 +65,11 @@ func (s *Store) processShard(ctx context.Context, shard *dynamodbstreams.Shard)
 		return
 	}
 
+	kafkaCfg := config.FromContext(ctx).Kafka
+	if kafkaCfg == nil {
+		logger.Fatal("Kafka config is nil")
+	}
+
 	shardIterator := iteratorOutput.ShardIterator
 	// Get records from shard iterator
 	for shardIterator != nil {
@@ -94,7 +100,7 @@ func (s *Store) processShard(ctx context.Context, shard *dynamodbstreams.Shard)
 				)
 			}
 
-			message, err := msg.KafkaMessage(ctx)
+			message, err := msg.KafkaMessage(*kafkaCfg)
 			if err != nil {
 				logger.Fatal("Failed to cast message from DynamoDB",
 					slog.Any("err", err),

sources/dynamodb/snapshot.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"log/slog"
 
+	"github.com/artie-labs/reader/config"
 	"github.com/artie-labs/reader/lib/dynamo"
 	"github.com/artie-labs/reader/lib/kafkalib"
 	"github.com/artie-labs/reader/lib/logger"
@@ -41,6 +42,11 @@ func (s *Store) streamAndPublish(ctx context.Context) error {
 		return fmt.Errorf("failed to retrieve primary keys, err: %v", err)
 	}
 
+	kafkaCfg := config.FromContext(ctx).Kafka
+	if kafkaCfg == nil {
+		return fmt.Errorf("kafka config is nil")
+	}
+
 	for _, file := range s.cfg.SnapshotSettings.SpecifiedFiles {
 		logFields := []any{
 			slog.String("fileName", *file.Key),
@@ -61,7 +67,7 @@ func (s *Store) streamAndPublish(ctx context.Context) error {
 				logger.Fatal("Failed to cast message from DynamoDB", slog.Any("err", err), slog.Any("msg", msg))
 			}
 
-			kafkaMsg, err := dynamoMsg.KafkaMessage(ctx)
+			kafkaMsg, err := dynamoMsg.KafkaMessage(*kafkaCfg)
 			if err != nil {
 				logger.Fatal("Failed to cast message from DynamoDB", slog.Any("err", err))
 			}