github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/storage/badger/collections.go (about)

     1  package badger
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	"github.com/dgraph-io/badger/v2"
     8  
     9  	"github.com/onflow/flow-go/model/flow"
    10  	"github.com/onflow/flow-go/storage"
    11  	"github.com/onflow/flow-go/storage/badger/operation"
    12  	"github.com/onflow/flow-go/storage/badger/transaction"
    13  )
    14  
    15  type Collections struct {
    16  	db           *badger.DB
    17  	transactions *Transactions
    18  }
    19  
    20  func NewCollections(db *badger.DB, transactions *Transactions) *Collections {
    21  	c := &Collections{
    22  		db:           db,
    23  		transactions: transactions,
    24  	}
    25  	return c
    26  }
    27  
    28  func (c *Collections) StoreLight(collection *flow.LightCollection) error {
    29  	err := operation.RetryOnConflict(c.db.Update, operation.InsertCollection(collection))
    30  	if err != nil {
    31  		return fmt.Errorf("could not insert collection: %w", err)
    32  	}
    33  
    34  	return nil
    35  }
    36  
    37  func (c *Collections) Store(collection *flow.Collection) error {
    38  	return operation.RetryOnConflictTx(c.db, transaction.Update, func(ttx *transaction.Tx) error {
    39  		light := collection.Light()
    40  		err := transaction.WithTx(operation.SkipDuplicates(operation.InsertCollection(&light)))(ttx)
    41  		if err != nil {
    42  			return fmt.Errorf("could not insert collection: %w", err)
    43  		}
    44  
    45  		for _, tx := range collection.Transactions {
    46  			err = c.transactions.storeTx(tx)(ttx)
    47  			if err != nil {
    48  				return fmt.Errorf("could not insert transaction: %w", err)
    49  			}
    50  		}
    51  
    52  		return nil
    53  	})
    54  }
    55  
    56  func (c *Collections) ByID(colID flow.Identifier) (*flow.Collection, error) {
    57  	var (
    58  		light      flow.LightCollection
    59  		collection flow.Collection
    60  	)
    61  
    62  	err := c.db.View(func(btx *badger.Txn) error {
    63  		err := operation.RetrieveCollection(colID, &light)(btx)
    64  		if err != nil {
    65  			return fmt.Errorf("could not retrieve collection: %w", err)
    66  		}
    67  
    68  		for _, txID := range light.Transactions {
    69  			tx, err := c.transactions.ByID(txID)
    70  			if err != nil {
    71  				return fmt.Errorf("could not retrieve transaction: %w", err)
    72  			}
    73  
    74  			collection.Transactions = append(collection.Transactions, tx)
    75  		}
    76  
    77  		return nil
    78  	})
    79  	if err != nil {
    80  		return nil, err
    81  	}
    82  
    83  	return &collection, nil
    84  }
    85  
    86  func (c *Collections) LightByID(colID flow.Identifier) (*flow.LightCollection, error) {
    87  	var collection flow.LightCollection
    88  
    89  	err := c.db.View(func(tx *badger.Txn) error {
    90  		err := operation.RetrieveCollection(colID, &collection)(tx)
    91  		if err != nil {
    92  			return fmt.Errorf("could not retrieve collection: %w", err)
    93  		}
    94  
    95  		return nil
    96  	})
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  
   101  	return &collection, nil
   102  }
   103  
   104  func (c *Collections) Remove(colID flow.Identifier) error {
   105  	return operation.RetryOnConflict(c.db.Update, func(btx *badger.Txn) error {
   106  		err := operation.RemoveCollection(colID)(btx)
   107  		if err != nil {
   108  			return fmt.Errorf("could not remove collection: %w", err)
   109  		}
   110  		return nil
   111  	})
   112  }
   113  
   114  func (c *Collections) StoreLightAndIndexByTransaction(collection *flow.LightCollection) error {
   115  	return operation.RetryOnConflict(c.db.Update, func(tx *badger.Txn) error {
   116  		err := operation.InsertCollection(collection)(tx)
   117  		if err != nil {
   118  			return fmt.Errorf("could not insert collection: %w", err)
   119  		}
   120  
   121  		for _, txID := range collection.Transactions {
   122  			err = operation.IndexCollectionByTransaction(txID, collection.ID())(tx)
   123  			if errors.Is(err, storage.ErrAlreadyExists) {
   124  				continue
   125  			}
   126  			if err != nil {
   127  				return fmt.Errorf("could not insert transaction ID: %w", err)
   128  			}
   129  		}
   130  
   131  		return nil
   132  	})
   133  }
   134  
   135  func (c *Collections) LightByTransactionID(txID flow.Identifier) (*flow.LightCollection, error) {
   136  	var collection flow.LightCollection
   137  	err := c.db.View(func(tx *badger.Txn) error {
   138  		collID := &flow.Identifier{}
   139  		err := operation.RetrieveCollectionID(txID, collID)(tx)
   140  		if err != nil {
   141  			return fmt.Errorf("could not retrieve collection id: %w", err)
   142  		}
   143  
   144  		err = operation.RetrieveCollection(*collID, &collection)(tx)
   145  		if err != nil {
   146  			return fmt.Errorf("could not retrieve collection: %w", err)
   147  		}
   148  
   149  		return nil
   150  	})
   151  	if err != nil {
   152  		return nil, err
   153  	}
   154  
   155  	return &collection, nil
   156  }