github.com/koko1123/flow-go-1@v0.29.6/storage/badger/operation/init.go (about)

     1  package operation
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  
     7  	"github.com/dgraph-io/badger/v3"
     8  
     9  	"github.com/koko1123/flow-go-1/storage"
    10  )
    11  
    12  // marker used to denote a database type
    13  type dbTypeMarker int
    14  
    15  func (marker dbTypeMarker) String() string {
    16  	return [...]string{
    17  		"dbMarkerPublic",
    18  		"dbMarkerSecret",
    19  	}[marker]
    20  }
    21  
    22  const (
    23  	// dbMarkerPublic denotes the public database
    24  	dbMarkerPublic dbTypeMarker = iota
    25  	// dbMarkerSecret denotes the secrets database
    26  	dbMarkerSecret
    27  )
    28  
    29  func InsertPublicDBMarker(txn *badger.Txn) error {
    30  	return insertDBTypeMarker(dbMarkerPublic)(txn)
    31  }
    32  
    33  func InsertSecretDBMarker(txn *badger.Txn) error {
    34  	return insertDBTypeMarker(dbMarkerSecret)(txn)
    35  }
    36  
    37  func EnsurePublicDB(db *badger.DB) error {
    38  	return ensureDBWithType(db, dbMarkerPublic)
    39  }
    40  
    41  func EnsureSecretDB(db *badger.DB) error {
    42  	return ensureDBWithType(db, dbMarkerSecret)
    43  }
    44  
    45  // insertDBTypeMarker inserts a database type marker if none exists. If a marker
    46  // already exists in the database, this function will return an error if the
    47  // marker does not match the argument, or return nil if it matches.
    48  func insertDBTypeMarker(marker dbTypeMarker) func(*badger.Txn) error {
    49  	return func(txn *badger.Txn) error {
    50  		var storedMarker dbTypeMarker
    51  		err := retrieveDBType(&storedMarker)(txn)
    52  		if err != nil && !errors.Is(err, storage.ErrNotFound) {
    53  			return fmt.Errorf("could not check db type marker: %w", err)
    54  		}
    55  
    56  		// we retrieved a marker from storage
    57  		if err == nil {
    58  			// the marker in storage does not match - error
    59  			if storedMarker != marker {
    60  				return fmt.Errorf("could not store db type marker - inconsistent marker already stored (expected: %s, actual: %s)", marker, storedMarker)
    61  			}
    62  			// the marker is already in storage - we're done
    63  			return nil
    64  		}
    65  
    66  		// no marker in storage, insert it
    67  		return insert(makePrefix(codeDBType), marker)(txn)
    68  	}
    69  }
    70  
    71  // ensureDBWithType ensures the given database has been initialized with the
    72  // given database type marker. If the given database has not been initialized
    73  // with any marker, or with a different marker than expected, returns an error.
    74  func ensureDBWithType(db *badger.DB, expectedMarker dbTypeMarker) error {
    75  	var actualMarker dbTypeMarker
    76  	err := db.View(retrieveDBType(&actualMarker))
    77  	if err != nil {
    78  		return fmt.Errorf("could not get db type: %w", err)
    79  	}
    80  	if actualMarker != expectedMarker {
    81  		return fmt.Errorf("wrong db type (expected: %s, actual: %s)", expectedMarker, actualMarker)
    82  	}
    83  	return nil
    84  }
    85  
    86  func retrieveDBType(marker *dbTypeMarker) func(*badger.Txn) error {
    87  	return retrieve(makePrefix(codeDBType), marker)
    88  }