0
Fork 0
mirror of https://github.com/project-zot/zot.git synced 2024-12-23 22:27:35 -05:00
zot/pkg/meta/dynamodb/iterator.go
LaurentiuNiculae 9e38ca51e3
feat(pagination): move pagination and sorting image summary results after conversion (#1637)
fix(config): check for config media type when pushing to repodb

Signed-off-by: Laurentiu Niculae <niculae.laurentiu1@gmail.com>
2023-07-31 22:16:09 +03:00

99 lines
2.3 KiB
Go

package dynamodb
import (
"context"
"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"
"zotregistry.io/zot/pkg/log"
)
type AttributesIterator interface {
First(ctx context.Context) (types.AttributeValue, error)
Next(ctx context.Context) (types.AttributeValue, error)
}
type BaseAttributesIterator struct {
Client *dynamodb.Client
Table string
Attribute string
itemBuffer []map[string]types.AttributeValue
currentItemIndex int
lastEvaluatedKey map[string]types.AttributeValue
readLimit *int32
log log.Logger
}
func NewBaseDynamoAttributesIterator(client *dynamodb.Client, table, attribute string, maxReadLimit int32,
log log.Logger,
) *BaseAttributesIterator {
var readLimit *int32
if maxReadLimit > 0 {
readLimit = &maxReadLimit
}
return &BaseAttributesIterator{
Client: client,
Table: table,
Attribute: attribute,
itemBuffer: []map[string]types.AttributeValue{},
currentItemIndex: 0,
readLimit: readLimit,
log: log,
}
}
func (dii *BaseAttributesIterator) First(ctx context.Context) (types.AttributeValue, error) {
scanOutput, err := dii.Client.Scan(ctx, &dynamodb.ScanInput{
TableName: aws.String(dii.Table),
Limit: dii.readLimit,
})
if err != nil {
return &types.AttributeValueMemberBOOL{}, err
}
if len(scanOutput.Items) == 0 {
return nil, nil
}
dii.itemBuffer = scanOutput.Items
dii.lastEvaluatedKey = scanOutput.LastEvaluatedKey
dii.currentItemIndex = 1
return dii.itemBuffer[0][dii.Attribute], nil
}
func (dii *BaseAttributesIterator) Next(ctx context.Context) (types.AttributeValue, error) {
if len(dii.itemBuffer) <= dii.currentItemIndex {
if dii.lastEvaluatedKey == nil {
return nil, nil
}
scanOutput, err := dii.Client.Scan(ctx, &dynamodb.ScanInput{
TableName: aws.String(dii.Table),
ExclusiveStartKey: dii.lastEvaluatedKey,
})
if err != nil {
return nil, err
}
// all items have been scanned
if len(scanOutput.Items) == 0 {
return nil, nil
}
dii.itemBuffer = scanOutput.Items
dii.lastEvaluatedKey = scanOutput.LastEvaluatedKey
dii.currentItemIndex = 0
}
nextItem := dii.itemBuffer[dii.currentItemIndex][dii.Attribute]
dii.currentItemIndex++
return nextItem, nil
}