diff --git a/internal/impl/aws/input_dynamodb_streams.go b/internal/impl/aws/input_dynamodb_streams.go new file mode 100644 index 0000000000..d6e0040eb4 --- /dev/null +++ b/internal/impl/aws/input_dynamodb_streams.go @@ -0,0 +1,912 @@ +package aws + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + "strings" + "sync" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/dynamodb" + "github.com/aws/aws-sdk-go-v2/service/dynamodbstreams" + "github.com/aws/aws-sdk-go-v2/service/dynamodbstreams/types" + "github.com/cenkalti/backoff/v4" + "github.com/gofrs/uuid" + + "github.com/warpstreamlabs/bento/internal/impl/aws/config" + "github.com/warpstreamlabs/bento/public/service" +) + +const ( + // DynamoDB Streams Input Fields + ddbsFieldTable = "table" + ddbsFieldStreamARN = "stream_arn" + ddbsFieldBatching = "batching" + ddbsFieldStartFromOldest = "start_from_oldest" + ddbsFieldCheckpointLimit = "checkpoint_limit" + ddbsFieldCommitPeriod = "commit_period" + ddbsFieldLeasePeriod = "lease_period" + ddbsFieldRebalancePeriod = "rebalance_period" + ddbsFieldPollInterval = "poll_interval" + + // DynamoDB Streams Checkpoint DynamoDB Fields + ddbsCPFieldTable = "table" + ddbsCPFieldCreate = "create" + ddbsCPFieldBillingMode = "billing_mode" + ddbsCPFieldReadCapacityUnits = "read_capacity_units" + ddbsCPFieldWriteCapacityUnits = "write_capacity_units" + ddbsFieldCheckpoint = "checkpoint" +) + +type ddbsConfig struct { + Table string + StreamARN string + StartFromOldest bool + CheckpointLimit int + CommitPeriod time.Duration + LeasePeriod time.Duration + RebalancePeriod time.Duration + PollInterval time.Duration + Checkpoint ddbsCheckpointConfig +} + +func ddbsInputConfigFromParsed(pConf *service.ParsedConfig) (conf ddbsConfig, err error) { + if conf.Table, err = pConf.FieldString(ddbsFieldTable); err != nil { + return + } + if conf.StreamARN, err = pConf.FieldString(ddbsFieldStreamARN); err != nil { + return + } + if conf.StartFromOldest, err = pConf.FieldBool(ddbsFieldStartFromOldest); err != nil { + return + } + if conf.CheckpointLimit, err = pConf.FieldInt(ddbsFieldCheckpointLimit); err != nil { + return + } + if conf.CommitPeriod, err = pConf.FieldDuration(ddbsFieldCommitPeriod); err != nil { + return + } + if conf.LeasePeriod, err = pConf.FieldDuration(ddbsFieldLeasePeriod); err != nil { + return + } + if conf.RebalancePeriod, err = pConf.FieldDuration(ddbsFieldRebalancePeriod); err != nil { + return + } + if conf.PollInterval, err = pConf.FieldDuration(ddbsFieldPollInterval); err != nil { + return + } + if pConf.Contains(ddbsFieldCheckpoint) { + if conf.Checkpoint, err = ddbsCheckpointConfigFromParsed(pConf.Namespace(ddbsFieldCheckpoint)); err != nil { + return + } + } + return +} + +func dynamoDBStreamsInputSpec() *service.ConfigSpec { + spec := service.NewConfigSpec(). + Beta(). + Categories("Services", "AWS"). + Summary("Consume messages from a DynamoDB stream."). + Description(` +Consumes change data capture events from a DynamoDB stream. Each message represents an INSERT, MODIFY, or REMOVE event on items in the source DynamoDB table. + +The latest consumed sequence for each shard is stored in a [DynamoDB table](#checkpoint), which allows this input to resume at the correct position during restarts. This table is also used for coordination across distributed instances when shard balancing. + +Bento will not store a consumed sequence unless it is acknowledged at the output level, which ensures at-least-once delivery guarantees. + +### Ordering + +By default messages of a shard can be processed in parallel, up to a limit determined by the field ` + "`checkpoint_limit`" + `. However, if strict ordered processing is required then this value must be set to 1 in order to process shard messages in lock-step. + +### Table Schema + +It's possible to configure Bento to create the DynamoDB checkpoint table if it does not already exist. However, if you wish to create this yourself (recommended) then create a table with a string HASH key ` + "`StreamID`" + ` and a string RANGE key ` + "`ShardID`" + `. + +### Stream ARN vs Table Name + +You can specify either a ` + "`table`" + ` name (in which case the latest stream ARN is resolved automatically) or an explicit ` + "`stream_arn`" + `. If both are specified, ` + "`stream_arn`" + ` takes precedence. + +### DynamoDB Streams Retention + +DynamoDB Streams retains data for 24 hours. If the consumer falls behind by more than 24 hours, the checkpoint will be stale and the shard iterator will be invalid. In this case the consumer will automatically reset to the oldest or latest position based on the ` + "`start_from_oldest`" + ` setting. +`).Fields( + service.NewStringField(ddbsFieldTable). + Description("The DynamoDB table name to consume the stream from. The stream ARN is resolved automatically. Either `table` or `stream_arn` must be specified."). + Default(""). + Examples("my-table"), + service.NewStringField(ddbsFieldStreamARN). + Description("An explicit DynamoDB stream ARN to consume from. Takes precedence over `table` if both are specified."). + Default(""). + Examples("arn:aws:dynamodb:us-east-1:123456789012:table/my-table/stream/2024-03-20T18:19:47.921"). + Advanced(), + service.NewObjectField(ddbsFieldCheckpoint, + service.NewStringField(ddbsCPFieldTable). + Description("The name of the DynamoDB table used for storing checkpoints."). + Default(""), + service.NewBoolField(ddbsCPFieldCreate). + Description("Whether, if the checkpoint table does not exist, it should be created."). + Default(false), + service.NewStringEnumField(ddbsCPFieldBillingMode, "PROVISIONED", "PAY_PER_REQUEST"). + Description("When creating the checkpoint table, determines the billing mode."). + Default("PAY_PER_REQUEST"). + Advanced(), + service.NewIntField(ddbsCPFieldReadCapacityUnits). + Description("Set the provisioned read capacity when creating the table with a `billing_mode` of `PROVISIONED`."). + Default(0). + Advanced(), + service.NewIntField(ddbsCPFieldWriteCapacityUnits). + Description("Set the provisioned write capacity when creating the table with a `billing_mode` of `PROVISIONED`."). + Default(0). + Advanced(), + ). + Description("Determines the table used for storing and accessing the latest consumed sequence for shards, and for coordinating balanced consumers of streams."), + service.NewIntField(ddbsFieldCheckpointLimit). + Description("The maximum gap between the in flight sequence versus the latest acknowledged sequence at a given time. Increasing this limit enables parallel processing and batching at the output level to work on individual shards. Any given sequence will not be committed unless all messages under that offset are delivered in order to preserve at least once delivery guarantees."). + Default(1024), + service.NewAutoRetryNacksToggleField(), + service.NewDurationField(ddbsFieldCommitPeriod). + Description("The period of time between each update to the checkpoint table."). + Default("5s"), + service.NewDurationField(ddbsFieldRebalancePeriod). + Description("The period of time between each attempt to rebalance shards across clients."). + Default("30s"). + Advanced(), + service.NewDurationField(ddbsFieldLeasePeriod). + Description("The period of time after which a client that has failed to update a shard checkpoint is assumed to be inactive."). + Default("30s"). + Advanced(), + service.NewBoolField(ddbsFieldStartFromOldest). + Description("Whether to consume from the oldest record when a sequence does not yet exist for the stream."). + Default(true), + service.NewDurationField(ddbsFieldPollInterval). + Description("The minimum interval between GetRecords calls for a given shard. DynamoDB Streams supports up to 5 GetRecords calls per second per shard."). + Default("1s"). + Advanced(), + ). + Fields(config.SessionFields()...). + Field(service.NewBatchPolicyField(ddbsFieldBatching)). + LintRule(`root = match { + this.table == "" && this.stream_arn == "" => ["either table or stream_arn must be specified"] +}`) + + return spec +} + +func init() { + err := service.RegisterBatchInput("aws_dynamodb_streams", dynamoDBStreamsInputSpec(), + func(conf *service.ParsedConfig, mgr *service.Resources) (service.BatchInput, error) { + r, err := newDynamoDBStreamsReaderFromParsed(conf, mgr) + if err != nil { + return nil, err + } + return service.AutoRetryNacksBatchedToggled(conf, r) + }) + if err != nil { + panic(err) + } +} + +//------------------------------------------------------------------------------ + +type dynamoDBStreamsReader struct { + conf ddbsConfig + clientID string + + sess aws.Config + batcher service.BatchPolicy + log *service.Logger + mgr *service.Resources + + boffPool sync.Pool + + streamsSvc *dynamodbstreams.Client + dynamodbSvc *dynamodb.Client + checkpointer *ddbsCheckpointer + + streamARN string + streamID string // Checkpoint key: table name if available, otherwise stream ARN. + + commitPeriod time.Duration + leasePeriod time.Duration + rebalancePeriod time.Duration + pollInterval time.Duration + + cMut sync.Mutex + msgChan chan asyncMessage + + ctx context.Context + done func() + + closeOnce sync.Once + closedChan chan struct{} +} + +func newDynamoDBStreamsReaderFromParsed(pConf *service.ParsedConfig, mgr *service.Resources) (*dynamoDBStreamsReader, error) { + conf, err := ddbsInputConfigFromParsed(pConf) + if err != nil { + return nil, err + } + sess, err := GetSession(context.TODO(), pConf) + if err != nil { + return nil, err + } + batcher, err := pConf.FieldBatchPolicy(ddbsFieldBatching) + if err != nil { + return nil, err + } + return newDynamoDBStreamsReaderFromConfig(conf, batcher, sess, mgr) +} + +func newDynamoDBStreamsReaderFromConfig(conf ddbsConfig, batcher service.BatchPolicy, sess aws.Config, mgr *service.Resources) (*dynamoDBStreamsReader, error) { + if batcher.IsNoop() { + batcher.Count = 1 + } + + r := dynamoDBStreamsReader{ + conf: conf, + sess: sess, + batcher: batcher, + log: mgr.Logger(), + mgr: mgr, + closedChan: make(chan struct{}), + commitPeriod: conf.CommitPeriod, + leasePeriod: conf.LeasePeriod, + rebalancePeriod: conf.RebalancePeriod, + pollInterval: conf.PollInterval, + } + r.ctx, r.done = context.WithCancel(context.Background()) + + u4, err := uuid.NewV4() + if err != nil { + return nil, err + } + r.clientID = u4.String() + + r.boffPool = sync.Pool{ + New: func() any { + boff := backoff.NewExponentialBackOff() + boff.InitialInterval = time.Millisecond * 300 + boff.MaxInterval = time.Second * 5 + boff.MaxElapsedTime = 0 + return boff + }, + } + + return &r, nil +} + +//------------------------------------------------------------------------------ + +// resolveStreamARN resolves the stream ARN either from the explicit config or +// by describing the DynamoDB table. +func (r *dynamoDBStreamsReader) resolveStreamARN(ctx context.Context) (string, error) { + if r.conf.StreamARN != "" { + return r.conf.StreamARN, nil + } + + out, err := r.dynamodbSvc.DescribeTable(ctx, &dynamodb.DescribeTableInput{ + TableName: aws.String(r.conf.Table), + }) + if err != nil { + return "", fmt.Errorf("failed to describe table %q: %w", r.conf.Table, err) + } + if out.Table.LatestStreamArn == nil || *out.Table.LatestStreamArn == "" { + return "", fmt.Errorf("table %q does not have streams enabled", r.conf.Table) + } + return *out.Table.LatestStreamArn, nil +} + +// collectShards returns all shards for the given stream ARN via pagination. +func (r *dynamoDBStreamsReader) collectShards(ctx context.Context, streamARN string) ([]types.Shard, error) { + var shards []types.Shard + var lastShardID *string + + for { + input := &dynamodbstreams.DescribeStreamInput{ + StreamArn: aws.String(streamARN), + ExclusiveStartShardId: lastShardID, + } + out, err := r.streamsSvc.DescribeStream(ctx, input) + if err != nil { + return nil, fmt.Errorf("failed to describe stream: %w", err) + } + if out.StreamDescription == nil { + return nil, errors.New("stream description is nil") + } + + shards = append(shards, out.StreamDescription.Shards...) + + lastShardID = out.StreamDescription.LastEvaluatedShardId + if lastShardID == nil { + break + } + } + + return shards, nil +} + +// isDDBShardOpen returns true if the shard does not have an ending sequence number, +// meaning it is still accepting writes. +func isDDBShardOpen(s types.Shard) bool { + if s.SequenceNumberRange == nil { + return true + } + return s.SequenceNumberRange.EndingSequenceNumber == nil +} + +func (r *dynamoDBStreamsReader) getShardIterator(ctx context.Context, shardID, sequence string) (string, error) { + iterType := types.ShardIteratorTypeTrimHorizon + if !r.conf.StartFromOldest { + iterType = types.ShardIteratorTypeLatest + } + + var startingSequence *string + if sequence != "" { + iterType = types.ShardIteratorTypeAfterSequenceNumber + startingSequence = &sequence + } + + res, err := r.streamsSvc.GetShardIterator(ctx, &dynamodbstreams.GetShardIteratorInput{ + StreamArn: aws.String(r.streamARN), + ShardId: aws.String(shardID), + ShardIteratorType: iterType, + SequenceNumber: startingSequence, + }) + if err != nil { + // If the sequence number is expired (>24h retention), fall back to + // the configured starting position. + var trimErr *types.TrimmedDataAccessException + if errors.As(err, &trimErr) && sequence != "" { + r.log.Warnf("Sequence expired for shard %v, resetting to %v", shardID, iterType) + return r.getShardIterator(ctx, shardID, "") + } + return "", err + } + + if res.ShardIterator == nil || *res.ShardIterator == "" { + return "", fmt.Errorf("failed to obtain shard iterator for shard %v", shardID) + } + return *res.ShardIterator, nil +} + +// getRecords fetches records from a shard iterator. Returns the records, the +// next iterator (empty string if shard is exhausted), and any error. On error +// the input iterator is always returned so callers can safely replace their +// stored iterator. +func (r *dynamoDBStreamsReader) getRecords(shardIter string) ([]types.Record, string, error) { + res, err := r.streamsSvc.GetRecords(r.ctx, &dynamodbstreams.GetRecordsInput{ + ShardIterator: aws.String(shardIter), + }) + if err != nil { + return nil, shardIter, err + } + + nextIter := "" + if res.NextShardIterator != nil { + nextIter = *res.NextShardIterator + } + return res.Records, nextIter, nil +} + +//------------------------------------------------------------------------------ + +type ddbsConsumerState int + +const ( + ddbsConsumerConsuming ddbsConsumerState = iota + ddbsConsumerYielding + ddbsConsumerFinished + ddbsConsumerClosing +) + +func (r *dynamoDBStreamsReader) runConsumer(wg *sync.WaitGroup, shardID, startingSequence string) (initErr error) { + defer func() { + if initErr != nil { + wg.Done() + if _, err := r.checkpointer.Checkpoint(context.Background(), r.streamID, shardID, startingSequence, true); err != nil { + r.log.Errorf("Failed to gracefully yield checkpoint: %v\n", err) + } + } + }() + + // Stores records, batches them up, and provides the batches for dispatch, + // whilst ensuring only N records are in flight at a given time. + var recordBatcher *ddbsRecordBatcher + if recordBatcher, initErr = r.newDDBSRecordBatcher(shardID, startingSequence); initErr != nil { + return initErr + } + + boff := r.boffPool.Get().(backoff.BackOff) + + var pending []types.Record + var iter string + if iter, initErr = r.getShardIterator(r.ctx, shardID, startingSequence); initErr != nil { + return initErr + } + + state := ddbsConsumerConsuming + var pendingMsg asyncMessage + + unblockedChan, blockedChan := make(chan time.Time), make(chan time.Time) + close(unblockedChan) + + var nextTimedBatchChan <-chan time.Time + var nextPullChan <-chan time.Time = unblockedChan + var nextFlushChan chan<- asyncMessage + commitCtx, commitCtxClose := context.WithTimeout(r.ctx, r.commitPeriod) + + go func() { + defer func() { + commitCtxClose() + recordBatcher.Close(context.Background(), state == ddbsConsumerFinished) + boff.Reset() + r.boffPool.Put(boff) + + reason := "" + switch state { + case ddbsConsumerFinished: + reason = " because the shard is closed" + if err := r.checkpointer.Complete(r.ctx, r.streamID, shardID, recordBatcher.GetSequence()); err != nil { + r.log.Errorf("Failed to mark shard '%v' as completed: %v", shardID, err) + } + case ddbsConsumerYielding: + reason = " because the shard has been claimed by another client" + if err := r.checkpointer.Yield(r.ctx, r.streamID, shardID, recordBatcher.GetSequence()); err != nil { + r.log.Errorf("Failed to yield checkpoint for stolen shard '%v': %v", shardID, err) + } + case ddbsConsumerClosing: + reason = " because the pipeline is shutting down" + if _, err := r.checkpointer.Checkpoint(context.Background(), r.streamID, shardID, recordBatcher.GetSequence(), true); err != nil { + r.log.Errorf("Failed to store final checkpoint for shard '%v': %v", shardID, err) + } + } + + wg.Done() + r.log.Debugf("Closing shard '%v' as client '%v'%v", shardID, r.checkpointer.clientID, reason) + }() + + r.log.Debugf("Consuming shard '%v' as client '%v'", shardID, r.checkpointer.clientID) + + unblockPullChan := func() { + if nextPullChan == blockedChan { + nextPullChan = unblockedChan + } + } + + for { + var err error + if state == ddbsConsumerConsuming && len(pending) == 0 && nextPullChan == unblockedChan { + if pending, iter, err = r.getRecords(iter); err != nil { + if !awsErrIsTimeout(err) { + nextPullChan = time.After(boff.NextBackOff()) + + var expiredErr *types.ExpiredIteratorException + var trimErr *types.TrimmedDataAccessException + if errors.As(err, &expiredErr) { + r.log.Warn("Shard iterator expired, attempting to refresh") + newIter, err := r.getShardIterator(r.ctx, shardID, recordBatcher.GetSequence()) + if err != nil { + r.log.Errorf("Failed to refresh shard iterator: %v", err) + } else { + iter = newIter + } + } else if errors.As(err, &trimErr) { + r.log.Warnf("Shard '%v' data trimmed (24h retention exceeded), resetting iterator", shardID) + newIter, err := r.getShardIterator(r.ctx, shardID, "") + if err != nil { + r.log.Errorf("Failed to refresh shard iterator after trim: %v", err) + } else { + iter = newIter + } + } else { + r.log.Errorf("Failed to pull DynamoDB stream records: %v\n", err) + } + } + } else if len(pending) == 0 { + // No records available, back off using poll interval + nextPullChan = time.After(r.pollInterval) + } else { + boff.Reset() + nextPullChan = blockedChan + } + if iter == "" { + state = ddbsConsumerFinished + } + } else { + unblockPullChan() + } + + if pendingMsg.msg == nil { + if len(pending) == 0 && state == ddbsConsumerFinished { + if pendingMsg, _ = recordBatcher.FlushMessage(r.ctx); pendingMsg.msg == nil { + return + } + } else if recordBatcher.HasPendingMessage() { + if pendingMsg, err = recordBatcher.FlushMessage(commitCtx); err != nil { + r.log.Errorf("Failed to dispatch message due to checkpoint error: %v\n", err) + } + } else if len(pending) > 0 { + var i int + var rec types.Record + for i, rec = range pending { + if recordBatcher.AddRecord(rec) { + if pendingMsg, err = recordBatcher.FlushMessage(commitCtx); err != nil { + r.log.Errorf("Failed to dispatch message due to checkpoint error: %v\n", err) + } + break + } + } + if pending = pending[i+1:]; len(pending) == 0 { + unblockPullChan() + } + } else { + unblockPullChan() + } + } + + if pendingMsg.msg != nil { + nextFlushChan = r.msgChan + } else { + nextFlushChan = nil + } + + if nextTimedBatchChan == nil { + if tNext, exists := recordBatcher.UntilNext(); exists { + nextTimedBatchChan = time.After(tNext) + } + } + + select { + case <-commitCtx.Done(): + if r.ctx.Err() != nil { + state = ddbsConsumerClosing + return + } + + commitCtxClose() + commitCtx, commitCtxClose = context.WithTimeout(r.ctx, r.commitPeriod) + + stillOwned, err := r.checkpointer.Checkpoint(r.ctx, r.streamID, shardID, recordBatcher.GetSequence(), false) + if err != nil { + r.log.Errorf("Failed to store checkpoint for shard '%v': %v", shardID, err) + } else if !stillOwned { + state = ddbsConsumerYielding + return + } + case <-nextTimedBatchChan: + nextTimedBatchChan = nil + case nextFlushChan <- pendingMsg: + pendingMsg = asyncMessage{} + case <-nextPullChan: + nextPullChan = unblockedChan + case <-r.ctx.Done(): + state = ddbsConsumerClosing + return + } + } + }() + return nil +} + +//------------------------------------------------------------------------------ + +func (r *dynamoDBStreamsReader) runBalancedShards() { + var wg sync.WaitGroup + defer func() { + wg.Wait() + r.closeOnce.Do(func() { + close(r.msgChan) + close(r.closedChan) + }) + }() + + for { + allShards, err := r.collectShards(r.ctx, r.streamARN) + var checkpointData *ddbsCheckpointData + if err == nil { + checkpointData, err = r.checkpointer.GetCheckpointsAndClaims(r.ctx, r.streamID) + } + if err != nil { + if r.ctx.Err() != nil { + return + } + r.log.Errorf("Failed to obtain stream shards or claims: %v", err) + select { + case <-time.After(r.rebalancePeriod): + continue + case <-r.ctx.Done(): + return + } + } + + clientClaims := checkpointData.ClientClaims + shardsWithCheckpoints := checkpointData.ShardsWithCheckpoints + completedShards := checkpointData.CompletedShards + + // Build parent map so we can enforce parent-before-child ordering. + parentOf := make(map[string]string, len(allShards)) + for _, s := range allShards { + if s.ParentShardId != nil && *s.ParentShardId != "" { + parentOf[*s.ShardId] = *s.ParentShardId + } + } + + unclaimedShards := make(map[string]string, len(allShards)) + for _, s := range allShards { + shardID := *s.ShardId + + // Skip completed shards — they're fully consumed. + if completedShards[shardID] { + continue + } + + // Skip child shards whose parent hasn't been fully consumed yet. + // This ensures per-key ordering is preserved across shard splits. + if parentID, hasParent := parentOf[shardID]; hasParent && !completedShards[parentID] { + continue + } + + if isDDBShardOpen(s) || shardsWithCheckpoints[shardID] { + unclaimedShards[shardID] = "" + } + } + for clientID, claims := range clientClaims { + for _, claim := range claims { + if time.Since(claim.LeaseTimeout) > r.leasePeriod*2 { + unclaimedShards[claim.ShardID] = clientID + } else { + delete(unclaimedShards, claim.ShardID) + } + } + } + + if len(unclaimedShards) > 0 { + for shardID, clientID := range unclaimedShards { + sequence, err := r.checkpointer.Claim(r.ctx, r.streamID, shardID, clientID) + if err != nil { + if r.ctx.Err() != nil { + return + } + if !errors.Is(err, ErrLeaseNotAcquired) { + r.log.Errorf("Failed to claim unclaimed shard '%v': %v", shardID, err) + } + continue + } + wg.Add(1) + if err = r.runConsumer(&wg, shardID, sequence); err != nil { + r.log.Errorf("Failed to start consumer: %v\n", err) + } + } + } else { + // No unclaimed shards — attempt to steal from overloaded clients. + selfClaims := len(clientClaims[r.clientID]) + for clientID, claims := range clientClaims { + if clientID == r.clientID { + continue + } + if len(claims) > (selfClaims + 1) { + targetShard := claims[rand.IntN(len(claims))].ShardID + r.log.Debugf("Attempting to steal shard '%v' from client '%v'", targetShard, clientID) + + sequence, err := r.checkpointer.Claim(r.ctx, r.streamID, targetShard, clientID) + if err != nil { + if r.ctx.Err() != nil { + return + } + if !errors.Is(err, ErrLeaseNotAcquired) { + r.log.Errorf("Failed to steal shard '%v': %v", targetShard, err) + } + continue + } + + wg.Add(1) + if err = r.runConsumer(&wg, targetShard, sequence); err != nil { + r.log.Errorf("Failed to start consumer: %v\n", err) + } else { + break + } + } + } + } + + select { + case <-time.After(r.rebalancePeriod): + case <-r.ctx.Done(): + return + } + } +} + +//------------------------------------------------------------------------------ + +// Connect establishes clients and starts background shard consumers. +func (r *dynamoDBStreamsReader) Connect(ctx context.Context) error { + r.cMut.Lock() + defer r.cMut.Unlock() + if r.msgChan != nil { + return nil + } + + r.streamsSvc = dynamodbstreams.NewFromConfig(r.sess) + r.dynamodbSvc = dynamodb.NewFromConfig(r.sess) + + streamARN, err := r.resolveStreamARN(ctx) + if err != nil { + return err + } + r.streamARN = streamARN + + // Use table name as checkpoint key when available — it's stable across + // stream re-creation (e.g., when streams are disabled and re-enabled on a + // table, a new ARN is generated but the table name stays the same). + r.streamID = r.streamARN + if r.conf.Table != "" { + r.streamID = r.conf.Table + } + + checkpointer, err := newDDBSCheckpointer(r.sess, r.clientID, r.conf.Checkpoint, r.leasePeriod, r.commitPeriod, r.log) + if err != nil { + return err + } + r.checkpointer = checkpointer + + // Verify the stream exists and is enabled. + out, err := r.streamsSvc.DescribeStream(ctx, &dynamodbstreams.DescribeStreamInput{ + StreamArn: aws.String(r.streamARN), + }) + if err != nil { + return fmt.Errorf("failed to describe stream: %w", err) + } + if out.StreamDescription != nil && out.StreamDescription.StreamStatus == types.StreamStatusDisabled { + return fmt.Errorf("stream %v is disabled", r.streamARN) + } + + r.msgChan = make(chan asyncMessage) + go r.runBalancedShards() + + r.log.Infof("Consuming DynamoDB stream: %v", r.streamARN) + return nil +} + +// ReadBatch attempts to read a message batch from the DynamoDB stream. +func (r *dynamoDBStreamsReader) ReadBatch(ctx context.Context) (service.MessageBatch, service.AckFunc, error) { + r.cMut.Lock() + msgChan := r.msgChan + r.cMut.Unlock() + + if msgChan == nil { + return nil, nil, service.ErrNotConnected + } + + select { + case m, open := <-msgChan: + if !open { + return nil, nil, service.ErrNotConnected + } + return m.msg, m.ackFn, nil + case <-ctx.Done(): + return nil, nil, ctx.Err() + } +} + +// Close gracefully shuts down the DynamoDB Streams input. +func (r *dynamoDBStreamsReader) Close(ctx context.Context) error { + r.done() + select { + case <-r.closedChan: + case <-ctx.Done(): + return ctx.Err() + } + return nil +} + +//------------------------------------------------------------------------------ + +// dynamoDBStreamEventToJSON converts a DynamoDB stream record into a structured +// JSON map suitable for message payloads. +func dynamoDBStreamEventToJSON(rec types.Record) map[string]any { + result := map[string]any{} + + if rec.EventID != nil { + result["eventID"] = *rec.EventID + } + if rec.EventName != "" { + result["eventName"] = string(rec.EventName) + } + if rec.EventVersion != nil { + result["eventVersion"] = *rec.EventVersion + } + if rec.EventSource != nil { + result["eventSource"] = *rec.EventSource + } + if rec.AwsRegion != nil { + result["awsRegion"] = *rec.AwsRegion + } + + if rec.Dynamodb != nil { + dynamo := map[string]any{} + + if rec.Dynamodb.Keys != nil { + dynamo["keys"] = streamsAttributeMapToJSON(rec.Dynamodb.Keys) + } + if rec.Dynamodb.NewImage != nil { + dynamo["newImage"] = streamsAttributeMapToJSON(rec.Dynamodb.NewImage) + } + if rec.Dynamodb.OldImage != nil { + dynamo["oldImage"] = streamsAttributeMapToJSON(rec.Dynamodb.OldImage) + } + if rec.Dynamodb.SequenceNumber != nil { + dynamo["sequenceNumber"] = *rec.Dynamodb.SequenceNumber + } + if rec.Dynamodb.SizeBytes != nil { + dynamo["sizeBytes"] = *rec.Dynamodb.SizeBytes + } + if rec.Dynamodb.StreamViewType != "" { + dynamo["streamViewType"] = string(rec.Dynamodb.StreamViewType) + } + if rec.Dynamodb.ApproximateCreationDateTime != nil { + dynamo["approximateCreationDateTime"] = rec.Dynamodb.ApproximateCreationDateTime.Format(time.RFC3339) + } + + result["dynamodb"] = dynamo + } + + return result +} + +// streamsAttributeMapToJSON converts a DynamoDB Streams attribute value map to +// a plain map[string]any representation. +func streamsAttributeMapToJSON(attrs map[string]types.AttributeValue) map[string]any { + result := make(map[string]any, len(attrs)) + for k, v := range attrs { + result[k] = streamsAttributeValueToJSON(v) + } + return result +} + +// streamsAttributeValueToJSON converts a single DynamoDB Streams attribute +// value to its Go native equivalent. +func streamsAttributeValueToJSON(av types.AttributeValue) any { + switch v := av.(type) { + case *types.AttributeValueMemberS: + return v.Value + case *types.AttributeValueMemberN: + return v.Value + case *types.AttributeValueMemberB: + return v.Value + case *types.AttributeValueMemberBOOL: + return v.Value + case *types.AttributeValueMemberNULL: + return nil + case *types.AttributeValueMemberSS: + return v.Value + case *types.AttributeValueMemberNS: + return v.Value + case *types.AttributeValueMemberBS: + return v.Value + case *types.AttributeValueMemberL: + list := make([]any, len(v.Value)) + for i, item := range v.Value { + list[i] = streamsAttributeValueToJSON(item) + } + return list + case *types.AttributeValueMemberM: + return streamsAttributeMapToJSON(v.Value) + default: + return nil + } +} + +// extractTableNameFromARN extracts the table name from a DynamoDB stream ARN. +// Format: arn:aws:dynamodb:region:account:table/TABLE_NAME/stream/TIMESTAMP +func extractTableNameFromARN(streamARN string) string { + parts := strings.Split(streamARN, "/") + if len(parts) >= 2 { + return parts[1] + } + return streamARN +} diff --git a/internal/impl/aws/input_dynamodb_streams_checkpointer.go b/internal/impl/aws/input_dynamodb_streams_checkpointer.go new file mode 100644 index 0000000000..e19ee72e4a --- /dev/null +++ b/internal/impl/aws/input_dynamodb_streams_checkpointer.go @@ -0,0 +1,414 @@ +package aws + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/dynamodb" + "github.com/aws/aws-sdk-go-v2/service/dynamodb/types" + + "github.com/warpstreamlabs/bento/public/service" +) + +type ddbsCheckpointConfig struct { + Table string + Create bool + BillingMode string + ReadCapacityUnits int64 + WriteCapacityUnits int64 +} + +func ddbsCheckpointConfigFromParsed(pConf *service.ParsedConfig) (conf ddbsCheckpointConfig, err error) { + if conf.Table, err = pConf.FieldString(ddbsCPFieldTable); err != nil { + return + } + if conf.Create, err = pConf.FieldBool(ddbsCPFieldCreate); err != nil { + return + } + if conf.BillingMode, err = pConf.FieldString(ddbsCPFieldBillingMode); err != nil { + return + } + if conf.ReadCapacityUnits, err = int64Field(pConf, ddbsCPFieldReadCapacityUnits); err != nil { + return + } + if conf.WriteCapacityUnits, err = int64Field(pConf, ddbsCPFieldWriteCapacityUnits); err != nil { + return + } + return +} + +const ddbsCompletedClientID = "_COMPLETED_" + +// ddbsCheckpointData extends the Kinesis checkpoint data with completed shard +// tracking, which is needed for parent-child shard ordering. +type ddbsCheckpointData struct { + awsKinesisCheckpointData + // CompletedShards is a set of shard IDs that have been fully consumed. + CompletedShards map[string]bool +} + +// ddbsCheckpointer manages shard checkpointing for the DynamoDB Streams input. +// It reuses the same DynamoDB table schema as the Kinesis checkpointer +// (StreamID + ShardID) for consistency. +type ddbsCheckpointer struct { + conf ddbsCheckpointConfig + clientID string + leaseDuration time.Duration + commitPeriod time.Duration + svc *dynamodb.Client + log *service.Logger +} + +func newDDBSCheckpointer( + aConf aws.Config, + clientID string, + conf ddbsCheckpointConfig, + leaseDuration time.Duration, + commitPeriod time.Duration, + log *service.Logger, +) (*ddbsCheckpointer, error) { + c := &ddbsCheckpointer{ + conf: conf, + leaseDuration: leaseDuration, + commitPeriod: commitPeriod, + svc: dynamodb.NewFromConfig(aConf), + clientID: clientID, + log: log, + } + + if err := c.ensureTableExists(context.TODO()); err != nil { + return nil, err + } + return c, nil +} + +func (c *ddbsCheckpointer) ensureTableExists(ctx context.Context) error { + _, err := c.svc.DescribeTable(ctx, &dynamodb.DescribeTableInput{ + TableName: aws.String(c.conf.Table), + }) + { + var aerr *types.ResourceNotFoundException + if err == nil || !errors.As(err, &aerr) { + return err + } + } + if !c.conf.Create { + return fmt.Errorf("checkpoint table %v does not exist", c.conf.Table) + } + + input := &dynamodb.CreateTableInput{ + AttributeDefinitions: []types.AttributeDefinition{ + {AttributeName: aws.String("StreamID"), AttributeType: types.ScalarAttributeTypeS}, + {AttributeName: aws.String("ShardID"), AttributeType: types.ScalarAttributeTypeS}, + }, + BillingMode: types.BillingMode(c.conf.BillingMode), + KeySchema: []types.KeySchemaElement{ + {AttributeName: aws.String("StreamID"), KeyType: types.KeyTypeHash}, + {AttributeName: aws.String("ShardID"), KeyType: types.KeyTypeRange}, + }, + TableName: aws.String(c.conf.Table), + } + if c.conf.BillingMode == "PROVISIONED" { + input.ProvisionedThroughput = &types.ProvisionedThroughput{ + ReadCapacityUnits: &c.conf.ReadCapacityUnits, + WriteCapacityUnits: &c.conf.WriteCapacityUnits, + } + } + if _, err = c.svc.CreateTable(ctx, input); err != nil { + return fmt.Errorf("failed to create checkpoint table: %w", err) + } + + // Wait for the table to become active before returning. Without this, + // the first checkpoint or claim operation would fail. + waiter := dynamodb.NewTableExistsWaiter(c.svc) + if err = waiter.Wait(ctx, &dynamodb.DescribeTableInput{ + TableName: aws.String(c.conf.Table), + }, 2*time.Minute); err != nil { + return fmt.Errorf("checkpoint table did not become active: %w", err) + } + + // Enable TTL on the ExpiresAt attribute so completed shard entries are + // automatically cleaned up after 48 hours, preventing unbounded growth. + if _, err = c.svc.UpdateTimeToLive(ctx, &dynamodb.UpdateTimeToLiveInput{ + TableName: aws.String(c.conf.Table), + TimeToLiveSpecification: &types.TimeToLiveSpecification{ + Enabled: aws.Bool(true), + AttributeName: aws.String("ExpiresAt"), + }, + }); err != nil { + // TTL enablement is best-effort — the input works without it, just + // accumulates completed shard entries over time. + c.log.Warnf("Failed to enable TTL on checkpoint table: %v", err) + } + + return nil +} + +// GetCheckpointsAndClaims retrieves all checkpoint data for a stream, including +// which shards have been marked as completed (fully consumed). +func (c *ddbsCheckpointer) GetCheckpointsAndClaims(ctx context.Context, streamID string) (*ddbsCheckpointData, error) { + result := &ddbsCheckpointData{ + awsKinesisCheckpointData: awsKinesisCheckpointData{ + ShardsWithCheckpoints: make(map[string]bool), + ClientClaims: make(map[string][]awsKinesisClientClaim), + }, + CompletedShards: make(map[string]bool), + } + + input := &dynamodb.QueryInput{ + TableName: aws.String(c.conf.Table), + KeyConditionExpression: aws.String("StreamID = :stream_id"), + ExpressionAttributeValues: map[string]types.AttributeValue{ + ":stream_id": &types.AttributeValueMemberS{ + Value: streamID, + }, + }, + } + + paginator := dynamodb.NewQueryPaginator(c.svc, input) + + for paginator.HasMorePages() { + page, err := paginator.NextPage(ctx) + if err != nil { + return nil, fmt.Errorf("failed to query checkpoints: %w", err) + } + + for _, item := range page.Items { + var shardID string + if s, ok := item["ShardID"].(*types.AttributeValueMemberS); ok { + shardID = s.Value + } + if shardID == "" { + continue + } + + result.ShardsWithCheckpoints[shardID] = true + + var clientID string + if s, ok := item["ClientID"].(*types.AttributeValueMemberS); ok { + clientID = s.Value + } + if clientID == "" { + continue + } + + // Shards marked with the completed sentinel are fully consumed. + if clientID == ddbsCompletedClientID { + result.CompletedShards[shardID] = true + continue + } + + var claim awsKinesisClientClaim + claim.ShardID = shardID + + if s, ok := item["LeaseTimeout"].(*types.AttributeValueMemberS); ok { + var parseErr error + if claim.LeaseTimeout, parseErr = time.Parse(time.RFC3339Nano, s.Value); parseErr != nil { + return nil, fmt.Errorf("failed to parse claim lease for shard %s: %w", shardID, parseErr) + } + } + if claim.LeaseTimeout.IsZero() { + return nil, fmt.Errorf("failed to extract lease timeout from claim for shard %s", shardID) + } + + result.ClientClaims[clientID] = append(result.ClientClaims[clientID], claim) + } + } + + return result, nil +} + +// Claim attempts to claim a shard. If fromClientID is specified the shard is +// stolen from that particular client. +func (c *ddbsCheckpointer) Claim(ctx context.Context, streamID, shardID, fromClientID string) (string, error) { + newLeaseTimeoutString := time.Now().Add(c.leaseDuration).Format(time.RFC3339Nano) + + var conditionalExpression string + expressionAttributeValues := map[string]types.AttributeValue{ + ":new_client_id": &types.AttributeValueMemberS{ + Value: c.clientID, + }, + ":new_lease_timeout": &types.AttributeValueMemberS{ + Value: newLeaseTimeoutString, + }, + } + + if fromClientID != "" { + conditionalExpression = "ClientID = :old_client_id" + expressionAttributeValues[":old_client_id"] = &types.AttributeValueMemberS{ + Value: fromClientID, + } + } else { + conditionalExpression = "attribute_not_exists(ClientID)" + } + + exp := "SET ClientID = :new_client_id, LeaseTimeout = :new_lease_timeout" + res, err := c.svc.UpdateItem(ctx, &dynamodb.UpdateItemInput{ + ReturnValues: types.ReturnValueAllOld, + TableName: &c.conf.Table, + ConditionExpression: &conditionalExpression, + UpdateExpression: &exp, + ExpressionAttributeValues: expressionAttributeValues, + Key: map[string]types.AttributeValue{ + "StreamID": &types.AttributeValueMemberS{Value: streamID}, + "ShardID": &types.AttributeValueMemberS{Value: shardID}, + }, + }) + if err != nil { + var aerr *types.ConditionalCheckFailedException + if errors.As(err, &aerr) { + return "", ErrLeaseNotAcquired + } + return "", err + } + + var startingSequence string + if s, ok := res.Attributes["SequenceNumber"].(*types.AttributeValueMemberS); ok { + startingSequence = s.Value + } + + var currentLease time.Time + if s, ok := res.Attributes["LeaseTimeout"].(*types.AttributeValueMemberS); ok { + currentLease, _ = time.Parse(time.RFC3339Nano, s.Value) + } + + // Wait a grace period when stealing from another client. + if fromClientID != "" && time.Since(currentLease) < c.leaseDuration { + waitFor := c.leaseDuration - time.Since(currentLease) + time.Second + select { + case <-time.After(waitFor): + case <-ctx.Done(): + return "", ctx.Err() + } + + cp, err := c.getCheckpoint(ctx, streamID, shardID) + if err != nil { + return "", err + } + if cp != nil { + startingSequence = cp.SequenceNumber + } + } + + return startingSequence, nil +} + +func (c *ddbsCheckpointer) getCheckpoint(ctx context.Context, streamID, shardID string) (*awsKinesisCheckpoint, error) { + rawItem, err := c.svc.GetItem(ctx, &dynamodb.GetItemInput{ + TableName: aws.String(c.conf.Table), + Key: map[string]types.AttributeValue{ + "ShardID": &types.AttributeValueMemberS{Value: shardID}, + "StreamID": &types.AttributeValueMemberS{Value: streamID}, + }, + }) + if err != nil { + var aerr *types.ResourceNotFoundException + if errors.As(err, &aerr) { + return nil, nil + } + return nil, err + } + + cp := awsKinesisCheckpoint{} + if s, ok := rawItem.Item["SequenceNumber"].(*types.AttributeValueMemberS); ok { + cp.SequenceNumber = s.Value + } else { + return nil, errors.New("sequence ID was not found in checkpoint") + } + if s, ok := rawItem.Item["ClientID"].(*types.AttributeValueMemberS); ok { + cp.ClientID = &s.Value + } + if s, ok := rawItem.Item["LeaseTimeout"].(*types.AttributeValueMemberS); ok { + timeout, err := time.Parse(time.RFC3339Nano, s.Value) + if err != nil { + return nil, err + } + cp.LeaseTimeout = &timeout + } + + return &cp, nil +} + +// Checkpoint sets a sequence number for a shard. Returns whether the shard is +// still owned by this client. +func (c *ddbsCheckpointer) Checkpoint(ctx context.Context, streamID, shardID, sequenceNumber string, final bool) (bool, error) { + item := map[string]types.AttributeValue{ + "StreamID": &types.AttributeValueMemberS{Value: streamID}, + "ShardID": &types.AttributeValueMemberS{Value: shardID}, + } + + if sequenceNumber != "" { + item["SequenceNumber"] = &types.AttributeValueMemberS{Value: sequenceNumber} + } + + if !final { + item["ClientID"] = &types.AttributeValueMemberS{Value: c.clientID} + item["LeaseTimeout"] = &types.AttributeValueMemberS{ + Value: time.Now().Add(c.leaseDuration).Format(time.RFC3339Nano), + } + } + + if _, err := c.svc.PutItem(ctx, &dynamodb.PutItemInput{ + ConditionExpression: aws.String("ClientID = :client_id"), + ExpressionAttributeValues: map[string]types.AttributeValue{ + ":client_id": &types.AttributeValueMemberS{Value: c.clientID}, + }, + TableName: aws.String(c.conf.Table), + Item: item, + }); err != nil { + var aerr *types.ConditionalCheckFailedException + if errors.As(err, &aerr) { + return false, nil + } + return false, err + } + return true, nil +} + +// Yield updates a checkpoint's sequence number without changing ownership. +func (c *ddbsCheckpointer) Yield(ctx context.Context, streamID, shardID, sequenceNumber string) error { + if sequenceNumber == "" { + return nil + } + + _, err := c.svc.UpdateItem(ctx, &dynamodb.UpdateItemInput{ + TableName: aws.String(c.conf.Table), + Key: map[string]types.AttributeValue{ + "StreamID": &types.AttributeValueMemberS{Value: streamID}, + "ShardID": &types.AttributeValueMemberS{Value: shardID}, + }, + ExpressionAttributeValues: map[string]types.AttributeValue{ + ":new_sequence_number": &types.AttributeValueMemberS{Value: sequenceNumber}, + }, + UpdateExpression: aws.String("SET SequenceNumber = :new_sequence_number"), + }) + return err +} + +// Complete marks a shard as fully consumed. The checkpoint entry is retained +// (rather than deleted) so that child shards can verify their parent completed +// before starting consumption, preserving per-key ordering across shard splits. +func (c *ddbsCheckpointer) Complete(ctx context.Context, streamID, shardID, sequenceNumber string) error { + item := map[string]types.AttributeValue{ + "StreamID": &types.AttributeValueMemberS{Value: streamID}, + "ShardID": &types.AttributeValueMemberS{Value: shardID}, + "ClientID": &types.AttributeValueMemberS{Value: ddbsCompletedClientID}, + // TTL: auto-delete after 48 hours. Completed shards only need to persist + // long enough for child shards to verify parent completion (DynamoDB Streams + // retains shard metadata for 24h, so 48h gives ample margin). + "ExpiresAt": &types.AttributeValueMemberN{ + Value: fmt.Sprintf("%d", time.Now().Add(48*time.Hour).Unix()), + }, + } + if sequenceNumber != "" { + item["SequenceNumber"] = &types.AttributeValueMemberS{Value: sequenceNumber} + } + _, err := c.svc.PutItem(ctx, &dynamodb.PutItemInput{ + TableName: aws.String(c.conf.Table), + Item: item, + }) + return err +} diff --git a/internal/impl/aws/input_dynamodb_streams_record_batcher.go b/internal/impl/aws/input_dynamodb_streams_record_batcher.go new file mode 100644 index 0000000000..4a15296158 --- /dev/null +++ b/internal/impl/aws/input_dynamodb_streams_record_batcher.go @@ -0,0 +1,138 @@ +package aws + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "time" + + "github.com/Jeffail/checkpoint" + "github.com/aws/aws-sdk-go-v2/service/dynamodbstreams/types" + + "github.com/warpstreamlabs/bento/public/service" +) + +type ddbsRecordBatcher struct { + streamARN string + shardID string + + batchPolicy *service.Batcher + checkpointer *checkpoint.Capped[string] + + flushedMessage service.MessageBatch + + batchedSequence string + + ackedSequence string + ackedMut sync.Mutex + ackedWG sync.WaitGroup +} + +func (r *dynamoDBStreamsReader) newDDBSRecordBatcher(shardID, sequence string) (*ddbsRecordBatcher, error) { + batchPolicy, err := r.batcher.NewBatcher(r.mgr) + if err != nil { + return nil, fmt.Errorf("failed to initialize batch policy for shard consumer: %w", err) + } + + return &ddbsRecordBatcher{ + streamARN: r.streamARN, + shardID: shardID, + batchPolicy: batchPolicy, + checkpointer: checkpoint.NewCapped[string](int64(r.conf.CheckpointLimit)), + ackedSequence: sequence, + }, nil +} + +// AddRecord converts a DynamoDB Streams record into a message and adds it to +// the current batch. Returns true if the batch is ready to flush. +func (b *ddbsRecordBatcher) AddRecord(rec types.Record) bool { + eventJSON := dynamoDBStreamEventToJSON(rec) + + data, err := json.Marshal(eventJSON) + if err != nil { + // Shouldn't happen for well-formed events, but handle defensively. + data = []byte("{}") + } + + p := service.NewMessage(data) + + // Set metadata for downstream processors/outputs. + tableName := extractTableNameFromARN(b.streamARN) + p.MetaSetMut("dynamodb_table", tableName) + p.MetaSetMut("dynamodb_stream_arn", b.streamARN) + p.MetaSetMut("dynamodb_shard_id", b.shardID) + + if rec.EventID != nil { + p.MetaSetMut("dynamodb_event_id", *rec.EventID) + } + if rec.EventName != "" { + p.MetaSetMut("dynamodb_event_name", string(rec.EventName)) + } + if rec.Dynamodb != nil && rec.Dynamodb.SequenceNumber != nil { + p.MetaSetMut("dynamodb_sequence_number", *rec.Dynamodb.SequenceNumber) + b.batchedSequence = *rec.Dynamodb.SequenceNumber + } + + if b.flushedMessage != nil { + b.flushedMessage = append(b.flushedMessage, p) + return true + } + return b.batchPolicy.Add(p) +} + +func (b *ddbsRecordBatcher) HasPendingMessage() bool { + return b.flushedMessage != nil +} + +func (b *ddbsRecordBatcher) FlushMessage(ctx context.Context) (asyncMessage, error) { + if b.flushedMessage == nil { + var err error + if b.flushedMessage, err = b.batchPolicy.Flush(ctx); err != nil || b.flushedMessage == nil { + return asyncMessage{}, err + } + } + + resolveFn, err := b.checkpointer.Track(ctx, b.batchedSequence, int64(len(b.flushedMessage))) + if err != nil { + if ctx.Err() != nil { + err = nil + } + return asyncMessage{}, err + } + + b.ackedWG.Add(1) + aMsg := asyncMessage{ + msg: b.flushedMessage, + ackFn: func(ctx context.Context, res error) error { + topSequence := resolveFn() + if topSequence != nil { + b.ackedMut.Lock() + b.ackedSequence = *topSequence + b.ackedMut.Unlock() + } + b.ackedWG.Done() + return res + }, + } + b.flushedMessage = nil + return aMsg, nil +} + +func (b *ddbsRecordBatcher) UntilNext() (time.Duration, bool) { + return b.batchPolicy.UntilNext() +} + +func (b *ddbsRecordBatcher) GetSequence() string { + b.ackedMut.Lock() + seq := b.ackedSequence + b.ackedMut.Unlock() + return seq +} + +func (b *ddbsRecordBatcher) Close(ctx context.Context, blocked bool) { + if blocked { + b.ackedWG.Wait() + } + _ = b.batchPolicy.Close(ctx) +} diff --git a/internal/impl/aws/input_dynamodb_streams_test.go b/internal/impl/aws/input_dynamodb_streams_test.go new file mode 100644 index 0000000000..f7539d0576 --- /dev/null +++ b/internal/impl/aws/input_dynamodb_streams_test.go @@ -0,0 +1,212 @@ +package aws + +import ( + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/service/dynamodbstreams/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestExtractTableNameFromARN(t *testing.T) { + t.Parallel() + tests := []struct { + name string + arn string + expected string + }{ + { + name: "standard stream ARN", + arn: "arn:aws:dynamodb:us-east-1:123456789012:table/my-table/stream/2024-03-20T18:19:47.921", + expected: "my-table", + }, + { + name: "table with hyphens", + arn: "arn:aws:dynamodb:eu-west-1:999999999999:table/events-prod-v2/stream/2025-01-01T00:00:00.000", + expected: "events-prod-v2", + }, + { + name: "no slashes fallback", + arn: "my-table", + expected: "my-table", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := extractTableNameFromARN(tt.arn) + assert.Equal(t, tt.expected, result) + }) + } +} + +func TestIsDDBShardOpen(t *testing.T) { + t.Parallel() + endSeq := "12345" + tests := []struct { + name string + shard types.Shard + expected bool + }{ + { + name: "nil sequence range", + shard: types.Shard{}, + expected: true, + }, + { + name: "no ending sequence", + shard: types.Shard{ + SequenceNumberRange: &types.SequenceNumberRange{ + StartingSequenceNumber: &endSeq, + }, + }, + expected: true, + }, + { + name: "has ending sequence (closed)", + shard: types.Shard{ + SequenceNumberRange: &types.SequenceNumberRange{ + StartingSequenceNumber: &endSeq, + EndingSequenceNumber: &endSeq, + }, + }, + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.expected, isDDBShardOpen(tt.shard)) + }) + } +} + +func TestDynamoDBStreamEventToJSON(t *testing.T) { + t.Parallel() + eventID := "event-123" + eventVersion := "1.1" + eventSource := "aws:dynamodb" + awsRegion := "us-east-1" + seqNum := "111111111111111111111111" + sizeBytes := int64(256) + creationTime := time.Date(2025, 1, 15, 10, 30, 0, 0, time.UTC) + + rec := types.Record{ + EventID: &eventID, + EventName: types.OperationTypeInsert, + EventVersion: &eventVersion, + EventSource: &eventSource, + AwsRegion: &awsRegion, + Dynamodb: &types.StreamRecord{ + Keys: map[string]types.AttributeValue{ + "PK": &types.AttributeValueMemberS{Value: "USER#123"}, + "SK": &types.AttributeValueMemberS{Value: "PROFILE"}, + }, + NewImage: map[string]types.AttributeValue{ + "PK": &types.AttributeValueMemberS{Value: "USER#123"}, + "SK": &types.AttributeValueMemberS{Value: "PROFILE"}, + "Name": &types.AttributeValueMemberS{Value: "Alice"}, + "Age": &types.AttributeValueMemberN{Value: "30"}, + }, + SequenceNumber: &seqNum, + SizeBytes: &sizeBytes, + StreamViewType: types.StreamViewTypeNewAndOldImages, + ApproximateCreationDateTime: &creationTime, + }, + } + + result := dynamoDBStreamEventToJSON(rec) + + assert.Equal(t, "event-123", result["eventID"]) + assert.Equal(t, "INSERT", result["eventName"]) + assert.Equal(t, "1.1", result["eventVersion"]) + assert.Equal(t, "aws:dynamodb", result["eventSource"]) + assert.Equal(t, "us-east-1", result["awsRegion"]) + + dynamo, ok := result["dynamodb"].(map[string]any) + require.True(t, ok, "dynamodb field should be a map") + + assert.Equal(t, seqNum, dynamo["sequenceNumber"]) + assert.Equal(t, int64(256), dynamo["sizeBytes"]) + assert.Equal(t, "NEW_AND_OLD_IMAGES", dynamo["streamViewType"]) + assert.Equal(t, "2025-01-15T10:30:00Z", dynamo["approximateCreationDateTime"]) + + keys, ok := dynamo["keys"].(map[string]any) + require.True(t, ok, "keys should be a map") + assert.Equal(t, "USER#123", keys["PK"]) + assert.Equal(t, "PROFILE", keys["SK"]) + + newImage, ok := dynamo["newImage"].(map[string]any) + require.True(t, ok, "newImage should be a map") + assert.Equal(t, "Alice", newImage["Name"]) + assert.Equal(t, "30", newImage["Age"]) + + assert.Nil(t, dynamo["oldImage"], "oldImage should not be present for INSERT") +} + +func TestAttributeValueToJSON(t *testing.T) { + tests := []struct { + name string + input types.AttributeValue + expected any + }{ + { + name: "string value", + input: &types.AttributeValueMemberS{Value: "hello"}, + expected: "hello", + }, + { + name: "number value", + input: &types.AttributeValueMemberN{Value: "42"}, + expected: "42", + }, + { + name: "bool value", + input: &types.AttributeValueMemberBOOL{Value: true}, + expected: true, + }, + { + name: "null value", + input: &types.AttributeValueMemberNULL{Value: true}, + expected: nil, + }, + { + name: "string set", + input: &types.AttributeValueMemberSS{Value: []string{"a", "b"}}, + expected: []string{"a", "b"}, + }, + { + name: "number set", + input: &types.AttributeValueMemberNS{Value: []string{"1", "2"}}, + expected: []string{"1", "2"}, + }, + { + name: "list value", + input: &types.AttributeValueMemberL{Value: []types.AttributeValue{ + &types.AttributeValueMemberS{Value: "item1"}, + &types.AttributeValueMemberN{Value: "99"}, + }}, + expected: []any{"item1", "99"}, + }, + { + name: "map value", + input: &types.AttributeValueMemberM{Value: map[string]types.AttributeValue{ + "nested": &types.AttributeValueMemberS{Value: "val"}, + }}, + expected: map[string]any{"nested": "val"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := streamsAttributeValueToJSON(tt.input) + assert.Equal(t, tt.expected, result) + }) + } +} + +func TestDynamoDBStreamsInputSpec(t *testing.T) { + spec := dynamoDBStreamsInputSpec() + require.NotNil(t, spec) +} diff --git a/internal/impl/aws/integration_dynamodb_streams_smoke_test.go b/internal/impl/aws/integration_dynamodb_streams_smoke_test.go new file mode 100644 index 0000000000..b3a607beec --- /dev/null +++ b/internal/impl/aws/integration_dynamodb_streams_smoke_test.go @@ -0,0 +1,201 @@ +package aws + +import ( + "context" + "fmt" + "net/http" + "strconv" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + awsconf "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/dynamodb" + dtypes "github.com/aws/aws-sdk-go-v2/service/dynamodb/types" + "github.com/aws/aws-sdk-go-v2/service/dynamodbstreams" + "github.com/ory/dockertest/v3" + "github.com/ory/dockertest/v3/docker" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/warpstreamlabs/bento/public/service/integration" +) + +func getLocalStackForStreams(t testing.TB) string { + t.Helper() + + portInt, err := integration.GetFreePort() + require.NoError(t, err) + + port := strconv.Itoa(portInt) + + pool, err := dockertest.NewPool("") + require.NoError(t, err) + pool.MaxWait = 5 * time.Minute + + resource, err := pool.RunWithOptions(&dockertest.RunOptions{ + Repository: "localstack/localstack", + Tag: "3.0", + ExposedPorts: []string{"4566/tcp"}, + PortBindings: map[docker.Port][]docker.PortBinding{ + docker.Port(port + "/tcp"): { + {HostIP: "", HostPort: port}, + }, + }, + Env: []string{ + "SERVICES=dynamodb,dynamodbstreams", + "LS_LOG=warn", + }, + }) + require.NoError(t, err) + port = resource.GetPort("4566/tcp") + t.Cleanup(func() { + assert.NoError(t, pool.Purge(resource)) + }) + _ = resource.Expire(900) + + require.NoError(t, pool.Retry(func() error { + resp, err := http.Get(fmt.Sprintf("http://localhost:%s/_localstack/health", port)) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("localstack not ready: %d", resp.StatusCode) + } + return nil + })) + + return port +} + +func TestIntegrationDynamoDBStreamsSmokeSDK(t *testing.T) { + integration.CheckSkip(t) + t.Parallel() + + port := getLocalStackForStreams(t) + endpoint := fmt.Sprintf("http://localhost:%v", port) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + conf, err := awsconf.LoadDefaultConfig(ctx, + awsconf.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("xxxxx", "xxxxx", "xxxxx")), + awsconf.WithRegion("us-east-1"), + ) + require.NoError(t, err) + conf.BaseEndpoint = &endpoint + + ddbClient := dynamodb.NewFromConfig(conf) + streamsClient := dynamodbstreams.NewFromConfig(conf) + + tableName := "smoke-test-streams" + + // Create table with streams enabled. + _, err = ddbClient.CreateTable(ctx, &dynamodb.CreateTableInput{ + TableName: aws.String(tableName), + AttributeDefinitions: []dtypes.AttributeDefinition{ + {AttributeName: aws.String("pk"), AttributeType: dtypes.ScalarAttributeTypeS}, + }, + KeySchema: []dtypes.KeySchemaElement{ + {AttributeName: aws.String("pk"), KeyType: dtypes.KeyTypeHash}, + }, + BillingMode: dtypes.BillingModePayPerRequest, + StreamSpecification: &dtypes.StreamSpecification{ + StreamEnabled: aws.Bool(true), + StreamViewType: dtypes.StreamViewTypeNewAndOldImages, + }, + }) + require.NoError(t, err) + + waiter := dynamodb.NewTableExistsWaiter(ddbClient) + require.NoError(t, waiter.Wait(ctx, &dynamodb.DescribeTableInput{ + TableName: aws.String(tableName), + }, time.Minute)) + + // Get stream ARN. + desc, err := ddbClient.DescribeTable(ctx, &dynamodb.DescribeTableInput{ + TableName: aws.String(tableName), + }) + require.NoError(t, err) + require.NotNil(t, desc.Table.LatestStreamArn) + streamARN := *desc.Table.LatestStreamArn + t.Logf("Stream ARN: %v", streamARN) + + // Describe stream to get shards. + streamDesc, err := streamsClient.DescribeStream(ctx, &dynamodbstreams.DescribeStreamInput{ + StreamArn: aws.String(streamARN), + }) + require.NoError(t, err) + require.NotNil(t, streamDesc.StreamDescription) + t.Logf("Shards: %d", len(streamDesc.StreamDescription.Shards)) + + for _, shard := range streamDesc.StreamDescription.Shards { + t.Logf(" Shard: %v, parent: %v", *shard.ShardId, shard.ParentShardId) + } + + require.NotEmpty(t, streamDesc.StreamDescription.Shards, "no shards found") + + // Get shard iterator. + shardID := *streamDesc.StreamDescription.Shards[0].ShardId + iterOut, err := streamsClient.GetShardIterator(ctx, &dynamodbstreams.GetShardIteratorInput{ + StreamArn: aws.String(streamARN), + ShardId: aws.String(shardID), + ShardIteratorType: "TRIM_HORIZON", + }) + require.NoError(t, err) + require.NotNil(t, iterOut.ShardIterator) + t.Log("Got shard iterator") + + // Write an item AFTER getting the iterator. + _, err = ddbClient.PutItem(ctx, &dynamodb.PutItemInput{ + TableName: aws.String(tableName), + Item: map[string]dtypes.AttributeValue{ + "pk": &dtypes.AttributeValueMemberS{Value: "test-1"}, + "data": &dtypes.AttributeValueMemberS{Value: "hello"}, + }, + }) + require.NoError(t, err) + t.Log("Item written after iterator") + + // Also try ListStreams to see what LocalStack reports. + listOut, err := streamsClient.ListStreams(ctx, &dynamodbstreams.ListStreamsInput{ + TableName: aws.String(tableName), + }) + require.NoError(t, err) + t.Logf("ListStreams returned %d streams", len(listOut.Streams)) + for _, s := range listOut.Streams { + t.Logf(" Stream: %v, table: %v", *s.StreamArn, *s.TableName) + } + + time.Sleep(2 * time.Second) // Give LocalStack time to propagate + + // Poll for records. + var records int + iter := *iterOut.ShardIterator + for i := 0; i < 10; i++ { + recOut, err := streamsClient.GetRecords(ctx, &dynamodbstreams.GetRecordsInput{ + ShardIterator: aws.String(iter), + }) + require.NoError(t, err) + records += len(recOut.Records) + t.Logf("Poll %d: got %d records (total %d)", i, len(recOut.Records), records) + + if records > 0 { + for _, r := range recOut.Records { + t.Logf(" Event: %v, ID: %v", r.EventName, *r.EventID) + } + break + } + + if recOut.NextShardIterator == nil { + t.Log("Shard exhausted") + break + } + iter = *recOut.NextShardIterator + time.Sleep(time.Second) + } + + require.Greater(t, records, 0, "expected at least one record from DynamoDB stream") +} diff --git a/internal/impl/aws/integration_dynamodb_streams_test.go b/internal/impl/aws/integration_dynamodb_streams_test.go new file mode 100644 index 0000000000..ead202e793 --- /dev/null +++ b/internal/impl/aws/integration_dynamodb_streams_test.go @@ -0,0 +1,319 @@ +package aws + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + awsconf "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/credentials" + "github.com/aws/aws-sdk-go-v2/service/dynamodb" + "github.com/aws/aws-sdk-go-v2/service/dynamodb/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/warpstreamlabs/bento/public/service" + "github.com/warpstreamlabs/bento/public/service/integration" + + _ "github.com/warpstreamlabs/bento/public/components/pure" +) + +// DynamoDB Streams requires LocalStack >= 3.0 (the shared GetLocalStack uses +// 4.9.2 which does not return stream records). The helper getLocalStackForStreams +// is defined in integration_dynamodb_streams_smoke_test.go. + +func createStreamEnabledTable(ctx context.Context, t testing.TB, port, tableName string) { + t.Helper() + + endpoint := fmt.Sprintf("http://localhost:%v", port) + conf, err := awsconf.LoadDefaultConfig(ctx, + awsconf.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("xxxxx", "xxxxx", "xxxxx")), + awsconf.WithRegion("us-east-1"), + ) + require.NoError(t, err) + + conf.BaseEndpoint = &endpoint + client := dynamodb.NewFromConfig(conf) + + _, err = client.CreateTable(ctx, &dynamodb.CreateTableInput{ + TableName: aws.String(tableName), + AttributeDefinitions: []types.AttributeDefinition{ + {AttributeName: aws.String("pk"), AttributeType: types.ScalarAttributeTypeS}, + }, + KeySchema: []types.KeySchemaElement{ + {AttributeName: aws.String("pk"), KeyType: types.KeyTypeHash}, + }, + BillingMode: types.BillingModePayPerRequest, + StreamSpecification: &types.StreamSpecification{ + StreamEnabled: aws.Bool(true), + StreamViewType: types.StreamViewTypeNewAndOldImages, + }, + }) + require.NoError(t, err) + + waiter := dynamodb.NewTableExistsWaiter(client) + require.NoError(t, waiter.Wait(ctx, &dynamodb.DescribeTableInput{ + TableName: aws.String(tableName), + }, time.Minute)) +} + +func putItems(ctx context.Context, t testing.TB, port, tableName string, items []map[string]types.AttributeValue) { + t.Helper() + + endpoint := fmt.Sprintf("http://localhost:%v", port) + conf, err := awsconf.LoadDefaultConfig(ctx, + awsconf.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("xxxxx", "xxxxx", "xxxxx")), + awsconf.WithRegion("us-east-1"), + ) + require.NoError(t, err) + + conf.BaseEndpoint = &endpoint + client := dynamodb.NewFromConfig(conf) + + for _, item := range items { + _, err := client.PutItem(ctx, &dynamodb.PutItemInput{ + TableName: aws.String(tableName), + Item: item, + }) + require.NoError(t, err) + } +} + +func ddbStreamsInputYAML(tableName, checkpointTable, port string) string { + return fmt.Sprintf(` +input: + aws_dynamodb_streams: + table: %v + start_from_oldest: true + checkpoint: + table: %v + create: true + commit_period: 1s + rebalance_period: 3s + lease_period: 5s + poll_interval: 500ms + endpoint: http://localhost:%v + region: us-east-1 + credentials: + id: xxxxx + secret: xxxxx + token: xxxxx + +output: + drop: {} +`, tableName, checkpointTable, port) +} + +func TestIntegrationDynamoDBStreams(t *testing.T) { + integration.CheckSkip(t) + t.Parallel() + + servicePort := getLocalStackForStreams(t) + + t.Run("basic_consumption", func(t *testing.T) { + testDynamoDBStreamsBasicConsumption(t, servicePort) + }) + + t.Run("checkpointing", func(t *testing.T) { + testDynamoDBStreamsCheckpointing(t, servicePort) + }) +} + +func testDynamoDBStreamsBasicConsumption(t *testing.T, port string) { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + tableName := "stream-basic-test" + checkpointTable := "stream-basic-test-cp" + + createStreamEnabledTable(ctx, t, port, tableName) + + builder := service.NewStreamBuilder() + require.NoError(t, builder.SetYAML(ddbStreamsInputYAML(tableName, checkpointTable, port))) + + var mu sync.Mutex + var received []map[string]any + var receivedMeta []map[string]string + require.NoError(t, builder.AddConsumerFunc(func(ctx context.Context, msg *service.Message) error { + b, err := msg.AsBytes() + if err != nil { + return err + } + var parsed map[string]any + if err := json.Unmarshal(b, &parsed); err != nil { + return err + } + + meta := make(map[string]string) + _ = msg.MetaWalk(func(key, value string) error { + meta[key] = value + return nil + }) + + mu.Lock() + received = append(received, parsed) + receivedMeta = append(receivedMeta, meta) + mu.Unlock() + return nil + })) + + stream, err := builder.Build() + require.NoError(t, err) + + streamCtx, streamCancel := context.WithCancel(ctx) + defer streamCancel() + errChan := make(chan error, 1) + go func() { + errChan <- stream.Run(streamCtx) + }() + + // Wait for the consumer to establish shard iterators. + time.Sleep(5 * time.Second) + + // Write items after consumer is running (LocalStack requires this ordering). + putItems(ctx, t, port, tableName, []map[string]types.AttributeValue{ + { + "pk": &types.AttributeValueMemberS{Value: "user-1"}, + "name": &types.AttributeValueMemberS{Value: "Alice"}, + }, + { + "pk": &types.AttributeValueMemberS{Value: "user-2"}, + "name": &types.AttributeValueMemberS{Value: "Bob"}, + }, + { + "pk": &types.AttributeValueMemberS{Value: "user-3"}, + "name": &types.AttributeValueMemberS{Value: "Charlie"}, + }, + }) + + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return len(received) >= 3 + }, 60*time.Second, 500*time.Millisecond, "timed out waiting for stream events") + + streamCancel() + <-errChan + + mu.Lock() + defer mu.Unlock() + + eventNames := make(map[string]bool) + pks := make(map[string]bool) + for i, evt := range received { + eventNames[evt["eventName"].(string)] = true + + dynamo, ok := evt["dynamodb"].(map[string]any) + require.True(t, ok, "event %d missing dynamodb field", i) + + newImage, ok := dynamo["newImage"].(map[string]any) + require.True(t, ok, "event %d missing newImage", i) + + pk, ok := newImage["pk"].(string) + require.True(t, ok, "event %d missing pk in newImage", i) + pks[pk] = true + + assert.Equal(t, tableName, receivedMeta[i]["dynamodb_table"]) + assert.NotEmpty(t, receivedMeta[i]["dynamodb_stream_arn"]) + assert.NotEmpty(t, receivedMeta[i]["dynamodb_shard_id"]) + assert.NotEmpty(t, receivedMeta[i]["dynamodb_event_id"]) + assert.Equal(t, "INSERT", receivedMeta[i]["dynamodb_event_name"]) + assert.NotEmpty(t, receivedMeta[i]["dynamodb_sequence_number"]) + } + + assert.True(t, eventNames["INSERT"], "expected INSERT events") + assert.True(t, pks["user-1"], "missing user-1") + assert.True(t, pks["user-2"], "missing user-2") + assert.True(t, pks["user-3"], "missing user-3") +} + +// testDynamoDBStreamsCheckpointing verifies that consuming records creates +// checkpoint entries in the DynamoDB checkpoint table with the expected schema. +func testDynamoDBStreamsCheckpointing(t *testing.T, port string) { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + + tableName := "stream-cp-test" + checkpointTable := "stream-cp-test-cp" + + createStreamEnabledTable(ctx, t, port, tableName) + + builder := service.NewStreamBuilder() + require.NoError(t, builder.SetYAML(ddbStreamsInputYAML(tableName, checkpointTable, port))) + + var mu sync.Mutex + var count int + require.NoError(t, builder.AddConsumerFunc(func(ctx context.Context, msg *service.Message) error { + mu.Lock() + count++ + mu.Unlock() + return nil + })) + + stream, err := builder.Build() + require.NoError(t, err) + + streamCtx, streamCancel := context.WithCancel(ctx) + defer streamCancel() + errChan := make(chan error, 1) + go func() { + errChan <- stream.Run(streamCtx) + }() + + time.Sleep(5 * time.Second) + + putItems(ctx, t, port, tableName, []map[string]types.AttributeValue{ + { + "pk": &types.AttributeValueMemberS{Value: "cp-item-1"}, + "data": &types.AttributeValueMemberS{Value: "one"}, + }, + }) + + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return count >= 1 + }, 30*time.Second, 500*time.Millisecond, "timed out waiting for events") + + // Wait for checkpoint to flush (commit_period is 1s). + time.Sleep(3 * time.Second) + + streamCancel() + <-errChan + + // Verify the checkpoint table was created and has entries. + endpoint := fmt.Sprintf("http://localhost:%v", port) + conf, err := awsconf.LoadDefaultConfig(ctx, + awsconf.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("xxxxx", "xxxxx", "xxxxx")), + awsconf.WithRegion("us-east-1"), + ) + require.NoError(t, err) + conf.BaseEndpoint = &endpoint + client := dynamodb.NewFromConfig(conf) + + scanOut, err := client.Scan(ctx, &dynamodb.ScanInput{ + TableName: aws.String(checkpointTable), + }) + require.NoError(t, err) + + require.NotEmpty(t, scanOut.Items, "checkpoint table should have at least one entry") + + for _, item := range scanOut.Items { + // Verify expected checkpoint schema fields. + streamID, ok := item["StreamID"].(*types.AttributeValueMemberS) + require.True(t, ok, "checkpoint missing StreamID") + assert.Equal(t, tableName, streamID.Value, "StreamID should be the table name") + + _, ok = item["ShardID"].(*types.AttributeValueMemberS) + require.True(t, ok, "checkpoint missing ShardID") + + _, ok = item["SequenceNumber"].(*types.AttributeValueMemberS) + assert.True(t, ok, "checkpoint should have SequenceNumber") + + t.Logf("Checkpoint entry: StreamID=%v, ShardID=%v", streamID.Value, item["ShardID"]) + } +}