badger_source.go raw

   1  //go:build !(js && wasm)
   2  
   3  package grapevine
   4  
   5  import (
   6  	"next.orly.dev/pkg/database"
   7  )
   8  
   9  // BadgerGraphSource adapts Badger's graph indexes to the GraphSource interface.
  10  type BadgerGraphSource struct {
  11  	db *database.D
  12  }
  13  
  14  // NewBadgerGraphSource creates a new BadgerGraphSource wrapping a Badger database.
  15  func NewBadgerGraphSource(db *database.D) *BadgerGraphSource {
  16  	return &BadgerGraphSource{db: db}
  17  }
  18  
  19  // TraverseFollowsOutbound does BFS outward from seed pubkey using the ppg index.
  20  func (s *BadgerGraphSource) TraverseFollowsOutbound(seedPubkey []byte, maxDepth int) (
  21  	pubkeysByDepth map[int][]string, allPubkeys map[string]int, err error,
  22  ) {
  23  	result, err := s.db.TraversePubkeyPubkey(seedPubkey, maxDepth, "out")
  24  	if err != nil {
  25  		return nil, nil, err
  26  	}
  27  	allPubkeys = make(map[string]int, result.TotalPubkeys)
  28  	for depth, pks := range result.PubkeysByDepth {
  29  		for _, pk := range pks {
  30  			allPubkeys[pk] = depth
  31  		}
  32  	}
  33  	return result.PubkeysByDepth, allPubkeys, nil
  34  }
  35  
  36  // GetFollowerPubkeys returns hex pubkeys of accounts that follow the target.
  37  func (s *BadgerGraphSource) GetFollowerPubkeys(targetHex string) ([]string, error) {
  38  	serial, err := s.db.PubkeyHexToSerial(targetHex)
  39  	if err != nil {
  40  		return nil, nil // unknown pubkey, no followers
  41  	}
  42  	followers, err := s.db.GetFollowersViaGPP(serial)
  43  	if err != nil {
  44  		return nil, err
  45  	}
  46  	result := make([]string, 0, len(followers))
  47  	for _, f := range followers {
  48  		h, err := s.db.GetPubkeyHexFromSerial(f)
  49  		if err != nil {
  50  			continue
  51  		}
  52  		result = append(result, h)
  53  	}
  54  	return result, nil
  55  }
  56  
  57  // GetFollowsPubkeys returns hex pubkeys that the source follows.
  58  func (s *BadgerGraphSource) GetFollowsPubkeys(sourceHex string) ([]string, error) {
  59  	serial, err := s.db.PubkeyHexToSerial(sourceHex)
  60  	if err != nil {
  61  		return nil, nil // unknown pubkey, no follows
  62  	}
  63  	follows, err := s.db.GetFollowsViaPPG(serial)
  64  	if err != nil {
  65  		return nil, err
  66  	}
  67  	result := make([]string, 0, len(follows))
  68  	for _, f := range follows {
  69  		h, err := s.db.GetPubkeyHexFromSerial(f)
  70  		if err != nil {
  71  			continue
  72  		}
  73  		result = append(result, h)
  74  	}
  75  	return result, nil
  76  }
  77