github.com/weaviate/weaviate@v1.24.6/adapters/repos/modules/modules.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 modulestorage
    13  
    14  import (
    15  	"fmt"
    16  	"os"
    17  
    18  	"github.com/pkg/errors"
    19  	"github.com/sirupsen/logrus"
    20  	"github.com/weaviate/weaviate/entities/moduletools"
    21  	bolt "go.etcd.io/bbolt"
    22  )
    23  
    24  type Repo struct {
    25  	logger  logrus.FieldLogger
    26  	baseDir string
    27  	db      *bolt.DB
    28  }
    29  
    30  func NewRepo(baseDir string, logger logrus.FieldLogger) (*Repo, error) {
    31  	r := &Repo{
    32  		baseDir: baseDir,
    33  		logger:  logger,
    34  	}
    35  
    36  	err := r.init()
    37  	return r, err
    38  }
    39  
    40  func (r *Repo) DBPath() string {
    41  	return fmt.Sprintf("%s/modules.db", r.baseDir)
    42  }
    43  
    44  func (r *Repo) DataPath() string {
    45  	return r.baseDir
    46  }
    47  
    48  func (r *Repo) init() error {
    49  	if err := os.MkdirAll(r.baseDir, 0o777); err != nil {
    50  		return errors.Wrapf(err, "create root path directory at %s", r.baseDir)
    51  	}
    52  
    53  	boltdb, err := bolt.Open(r.DBPath(), 0o600, nil)
    54  	if err != nil {
    55  		return errors.Wrapf(err, "open bolt at %s", r.DBPath())
    56  	}
    57  
    58  	r.db = boltdb
    59  
    60  	return nil
    61  }
    62  
    63  type storageBucket struct {
    64  	bucketKey []byte
    65  	repo      *Repo
    66  }
    67  
    68  func (r *Repo) Storage(bucketName string) (moduletools.Storage, error) {
    69  	storage := &storageBucket{
    70  		bucketKey: []byte(bucketName),
    71  		repo:      r,
    72  	}
    73  
    74  	err := storage.init()
    75  	return storage, err
    76  }
    77  
    78  func (s *storageBucket) init() error {
    79  	return s.repo.db.Update(func(tx *bolt.Tx) error {
    80  		if _, err := tx.CreateBucketIfNotExists(s.bucketKey); err != nil {
    81  			return errors.Wrapf(err, "create module storage bucket '%s'",
    82  				string(s.bucketKey))
    83  		}
    84  
    85  		return nil
    86  	})
    87  }
    88  
    89  func (s *storageBucket) Put(key, value []byte) error {
    90  	return s.repo.db.Batch(func(tx *bolt.Tx) error {
    91  		b := tx.Bucket(s.bucketKey)
    92  		if b == nil {
    93  			return errors.Errorf("no bucket for key %s found", string(s.bucketKey))
    94  		}
    95  
    96  		if err := b.Put(key, value); err != nil {
    97  			return errors.Wrapf(err, "put value for key %s", string(key))
    98  		}
    99  
   100  		return nil
   101  	})
   102  }
   103  
   104  func (s *storageBucket) Get(key []byte) ([]byte, error) {
   105  	var out []byte
   106  	err := s.repo.db.View(func(tx *bolt.Tx) error {
   107  		b := tx.Bucket(s.bucketKey)
   108  		if b == nil {
   109  			return errors.Errorf("no bucket for key %s found", string(s.bucketKey))
   110  		}
   111  
   112  		out = b.Get(key)
   113  		return nil
   114  	})
   115  
   116  	return out, err
   117  }
   118  
   119  func (s *storageBucket) Scan(scan moduletools.ScanFn) error {
   120  	err := s.repo.db.View(func(tx *bolt.Tx) error {
   121  		b := tx.Bucket(s.bucketKey)
   122  
   123  		c := b.Cursor()
   124  		for k, v := c.First(); k != nil; k, v = c.Next() {
   125  			ok, err := scan(k, v)
   126  			if err != nil {
   127  				return errors.Wrapf(err, "read item %q", string(k))
   128  			}
   129  
   130  			if !ok {
   131  				break
   132  			}
   133  		}
   134  
   135  		return nil
   136  	})
   137  
   138  	return err
   139  }