github.com/weaviate/weaviate@v1.24.6/adapters/repos/transactions/store.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package txstore
    13  
    14  import (
    15  	"context"
    16  	"encoding/json"
    17  	"fmt"
    18  	"os"
    19  	"path"
    20  
    21  	"github.com/sirupsen/logrus"
    22  	"github.com/weaviate/weaviate/usecases/cluster"
    23  	"go.etcd.io/bbolt"
    24  )
    25  
    26  var txBucket = []byte("transactions")
    27  
    28  type Store struct {
    29  	db           *bbolt.DB
    30  	log          logrus.FieldLogger
    31  	homeDir      string
    32  	unmarshaller unmarshalFn
    33  }
    34  
    35  func NewStore(homeDir string, logger logrus.FieldLogger) *Store {
    36  	return &Store{
    37  		homeDir: homeDir,
    38  		log:     logger,
    39  	}
    40  }
    41  
    42  func (s *Store) SetUmarshalFn(fn unmarshalFn) {
    43  	s.unmarshaller = fn
    44  }
    45  
    46  func (s *Store) Open() error {
    47  	if err := os.MkdirAll(s.homeDir, 0o777); err != nil {
    48  		return fmt.Errorf("create root directory %q: %w", s.homeDir, err)
    49  	}
    50  
    51  	path := path.Join(s.homeDir, "tx.db")
    52  	boltDB, err := initBoltDB(path)
    53  	if err != nil {
    54  		return fmt.Errorf("init bolt_db: %w", err)
    55  	}
    56  
    57  	s.db = boltDB
    58  
    59  	return nil
    60  }
    61  
    62  func (s *Store) StoreTx(ctx context.Context, tx *cluster.Transaction) error {
    63  	data, err := json.Marshal(txWrapper{
    64  		ID:      tx.ID,
    65  		Payload: tx.Payload,
    66  		Type:    tx.Type,
    67  	})
    68  	if err != nil {
    69  		return fmt.Errorf("marshal tx: %w", err)
    70  	}
    71  
    72  	return s.db.Update(func(boltTx *bbolt.Tx) error {
    73  		b := boltTx.Bucket(txBucket)
    74  		return b.Put([]byte(tx.ID), data)
    75  	})
    76  }
    77  
    78  func (s *Store) DeleteTx(ctx context.Context, txId string) error {
    79  	return s.db.Update(func(boltTx *bbolt.Tx) error {
    80  		b := boltTx.Bucket(txBucket)
    81  		return b.Delete([]byte(txId))
    82  	})
    83  }
    84  
    85  func (s *Store) IterateAll(ctx context.Context,
    86  	cb func(tx *cluster.Transaction),
    87  ) error {
    88  	return s.db.View(func(boltTx *bbolt.Tx) error {
    89  		b := boltTx.Bucket(txBucket)
    90  		c := b.Cursor()
    91  		for k, v := c.First(); k != nil; k, v = c.Next() {
    92  			var txWrap txWrapperRead
    93  			if err := json.Unmarshal(v, &txWrap); err != nil {
    94  				return err
    95  			}
    96  
    97  			tx := cluster.Transaction{
    98  				ID:   txWrap.ID,
    99  				Type: txWrap.Type,
   100  			}
   101  
   102  			pl, err := s.unmarshaller(tx.Type, txWrap.Payload)
   103  			if err != nil {
   104  				return err
   105  			}
   106  
   107  			tx.Payload = pl
   108  
   109  			cb(&tx)
   110  
   111  		}
   112  		return nil
   113  	})
   114  }
   115  
   116  func (s *Store) Close() error {
   117  	return nil
   118  }
   119  
   120  func initBoltDB(filePath string) (*bbolt.DB, error) {
   121  	db, err := bbolt.Open(filePath, 0o600, nil)
   122  	if err != nil {
   123  		return nil, fmt.Errorf("open %q: %w", filePath, err)
   124  	}
   125  
   126  	root := func(tx *bbolt.Tx) error {
   127  		_, err := tx.CreateBucketIfNotExists(txBucket)
   128  		return err
   129  	}
   130  
   131  	return db, db.Update(root)
   132  }
   133  
   134  type txWrapper struct {
   135  	ID      string                  `json:"id"`
   136  	Payload any                     `json:"payload"`
   137  	Type    cluster.TransactionType `json:"type"`
   138  }
   139  
   140  // delayed unmarshalling of the payload, so we can inject a specific
   141  // marshaller
   142  type txWrapperRead struct {
   143  	ID      string                  `json:"id"`
   144  	Payload json.RawMessage         `json:"payload"`
   145  	Type    cluster.TransactionType `json:"type"`
   146  }
   147  
   148  type unmarshalFn func(txType cluster.TransactionType, payload json.RawMessage) (any, error)