zotregistry.io/zot@v1.4.4-0.20231124084042-02a8ed785457/pkg/meta/dynamodb/iterator.go (about) 1 package dynamodb 2 3 import ( 4 "context" 5 6 "github.com/aws/aws-sdk-go-v2/aws" 7 "github.com/aws/aws-sdk-go-v2/service/dynamodb" 8 "github.com/aws/aws-sdk-go-v2/service/dynamodb/types" 9 10 "zotregistry.io/zot/pkg/log" 11 ) 12 13 type AttributesIterator interface { 14 First(ctx context.Context) (types.AttributeValue, error) 15 Next(ctx context.Context) (types.AttributeValue, error) 16 } 17 18 type BaseAttributesIterator struct { 19 Client *dynamodb.Client 20 Table string 21 Attribute string 22 23 itemBuffer []map[string]types.AttributeValue 24 currentItemIndex int 25 lastEvaluatedKey map[string]types.AttributeValue 26 readLimit *int32 27 28 log log.Logger 29 } 30 31 func NewBaseDynamoAttributesIterator(client *dynamodb.Client, table, attribute string, maxReadLimit int32, 32 log log.Logger, 33 ) *BaseAttributesIterator { 34 var readLimit *int32 35 36 if maxReadLimit > 0 { 37 readLimit = &maxReadLimit 38 } 39 40 return &BaseAttributesIterator{ 41 Client: client, 42 Table: table, 43 Attribute: attribute, 44 itemBuffer: []map[string]types.AttributeValue{}, 45 currentItemIndex: 0, 46 readLimit: readLimit, 47 log: log, 48 } 49 } 50 51 func (dii *BaseAttributesIterator) First(ctx context.Context) (types.AttributeValue, error) { 52 scanOutput, err := dii.Client.Scan(ctx, &dynamodb.ScanInput{ 53 TableName: aws.String(dii.Table), 54 Limit: dii.readLimit, 55 ProjectionExpression: aws.String(dii.Attribute), 56 }) 57 if err != nil { 58 return &types.AttributeValueMemberBOOL{}, err 59 } 60 61 if len(scanOutput.Items) == 0 { 62 return nil, nil 63 } 64 65 dii.itemBuffer = scanOutput.Items 66 dii.lastEvaluatedKey = scanOutput.LastEvaluatedKey 67 dii.currentItemIndex = 1 68 69 return dii.itemBuffer[0][dii.Attribute], nil 70 } 71 72 func (dii *BaseAttributesIterator) Next(ctx context.Context) (types.AttributeValue, error) { 73 if len(dii.itemBuffer) <= dii.currentItemIndex { 74 if dii.lastEvaluatedKey == nil { 75 return nil, nil 76 } 77 78 scanOutput, err := dii.Client.Scan(ctx, &dynamodb.ScanInput{ 79 TableName: aws.String(dii.Table), 80 ExclusiveStartKey: dii.lastEvaluatedKey, 81 }) 82 if err != nil { 83 return nil, err 84 } 85 86 // all items have been scanned 87 if len(scanOutput.Items) == 0 { 88 return nil, nil 89 } 90 91 dii.itemBuffer = scanOutput.Items 92 dii.lastEvaluatedKey = scanOutput.LastEvaluatedKey 93 dii.currentItemIndex = 0 94 } 95 96 nextItem := dii.itemBuffer[dii.currentItemIndex][dii.Attribute] 97 dii.currentItemIndex++ 98 99 return nextItem, nil 100 }