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  }