dndex/storage/driver/boltdb.go

package driver

import (
	"context"
	"errors"
	"fmt"
	"local/dndex/config"
	"local/dndex/storage/entity"
	"os"
	"regexp"

	"github.com/boltdb/bolt"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

var (
	errModifiedReserved = errors.New("cannot modify reserved field " + entity.Name)
)

type BoltDB struct {
	db *bolt.DB
}

func NewBoltDB() *BoltDB {
	config := config.New()
	db, err := bolt.Open(config.DBURI, os.ModePerm, nil)
	if err != nil {
		panic(err)
	}
	return &BoltDB{
		db: db,
	}
}

func (bdb *BoltDB) count(ctx context.Context, namespace string, filter interface{}) (int, error) {
	ch, err := bdb.Find(ctx, namespace, filter)
	n := 0
	for _ = range ch {
		n++
	}
	return n, err
}

func (bdb *BoltDB) Find(_ context.Context, namespace string, filter interface{}) (chan bson.Raw, error) {
	b, err := bson.Marshal(filter)
	if err != nil {
		return nil, err
	}
	m := bson.M{}
	if err := bson.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	results := make([]bson.Raw, 0)
	err = bdb.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(namespace))
		if bucket == nil {
			return nil
		}
		cursor := bucket.Cursor()
		for k, v := cursor.First(); k != nil && v != nil; k, v = cursor.Next() {
			n := bson.M{}
			if err := bson.Unmarshal(v, &n); err != nil {
				return err
			}
			if matches(n, m) {
				results = append(results, bson.Raw(v))
			}
		}
		return nil
	})
	ch := make(chan bson.Raw)
	go func() {
		defer close(ch)
		for i := range results {
			ch <- results[i]
		}
	}()
	return ch, err
}

func (bdb *BoltDB) Update(ctx context.Context, namespace string, filter, operator interface{}) error {
	b, err := bson.Marshal(filter)
	if err != nil {
		return err
	}
	m := bson.M{}
	if err := bson.Unmarshal(b, &m); err != nil {
		return err
	}
	return bdb.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(namespace))
		if bucket == nil {
			return nil
		}
		cursor := bucket.Cursor()
		for k, v := cursor.First(); k != nil && v != nil; k, v = cursor.Next() {
			n := bson.M{}
			if err := bson.Unmarshal(v, &n); err != nil {
				return err
			}
			if matches(n, m) {
				n, err := apply(n, operator)
				if err != nil {
					return err
				}
				v, err := bson.Marshal(n)
				if err != nil {
					return err
				}
				if err := bucket.Put(k, v); err != nil {
					return err
				}
			}
		}
		return nil
	})
}

func (bdb *BoltDB) Insert(ctx context.Context, namespace string, doc interface{}) error {
	b, err := bson.Marshal(doc)
	if err != nil {
		return err
	}
	m := bson.M{}
	if err := bson.Unmarshal(b, &m); err != nil {
		return err
	}
	if _, ok := m[entity.Name]; !ok {
		return errors.New("primary key required to insert: did not find " + entity.Name)
	} else if _, ok := m[entity.Name].(string); !ok {
		return errors.New("primary key must be a string")
	}
	return bdb.db.Update(func(tx *bolt.Tx) error {
		bucket, err := tx.CreateBucketIfNotExists([]byte(namespace))
		if err != nil {
			return err
		}
		k := []byte(m[entity.Name].(string))
		v := bucket.Get(k)
		if len(v) > 0 {
			return errors.New("cannot insert: collision on primary key")
		}
		return bucket.Put(k, b)
	})
}

func (bdb *BoltDB) Delete(_ context.Context, namespace string, filter interface{}) error {
	b, err := bson.Marshal(filter)
	if err != nil {
		return err
	}
	m := bson.M{}
	if err := bson.Unmarshal(b, &m); err != nil {
		return err
	}
	return bdb.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(namespace))
		if bucket == nil {
			return nil
		}
		cursor := bucket.Cursor()
		for k, v := cursor.First(); k != nil && v != nil; k, v = cursor.Next() {
			n := bson.M{}
			if err := bson.Unmarshal(v, &n); err != nil {
				return err
			}
			if matches(n, m) {
				if err := bucket.Delete(k); err != nil {
					return err
				}
			}
		}
		return nil
	})
}

func matches(doc, filter bson.M) bool {
	for k, v := range filter {
		if _, ok := doc[k]; !ok {
			continue
		}
		switch v.(type) {
		case map[string]interface{}, primitive.M:
			m, ok := v.(map[string]interface{})
			if !ok {
				m = map[string]interface{}(v.(primitive.M))
			}
			for k2, v2 := range m {
				switch k2 {
				case "$regex":
					pattern, ok := v2.(string)
					if !ok {
						return false
					}
					re, err := regexp.Compile(pattern)
					if err != nil {
						return false
					}
					if !re.MatchString(fmt.Sprint(doc[k])) {
						return false
					}
				case "$in":
					options, ok := v2.([]interface{})
					if !ok {
						options = []interface{}(v2.(primitive.A))
						ok = true
					}
					matches := false
					for _, option := range options {
						if fmt.Sprint(doc[k]) == fmt.Sprint(option) {
							matches = true
						}
					}
					if !matches {
						return false
					}
				default:
					dock, ok := doc[k].(map[string]interface{})
					if !ok {
						return false
					}
					if !matches(dock, bson.M(m)) {
						return false
					}
				}
			}
		default:
			if fmt.Sprint(v) != fmt.Sprint(doc[k]) {
				return false
			}
		}
	}
	return true
}

func apply(doc bson.M, operator interface{}) (bson.M, error) {
	b, err := bson.Marshal(operator)
	if err != nil {
		return nil, err
	}
	op := bson.M{}
	if err := bson.Unmarshal(b, &op); err != nil {
		return nil, err
	}
	if _, ok := op[entity.Name]; ok {
		return nil, errModifiedReserved
	}
	for k, v := range op {
		operateBM, ok := v.(bson.M)
		if !ok {
			operateM, mok := v.(map[string]interface{})
			ok = mok
			if ok {
				operateBM = bson.M(operateM)
			}
		}
		if !ok {
			return nil, fmt.Errorf("invalid apply operand: %s:%T", k, v)
		}
		var op func(bson.M, bson.M) (bson.M, error)
		switch k {
		case "$set":
			op = applySet
		default:
			return nil, errors.New("cannot apply operation " + k)
		}
		doc, err = op(doc, operateBM)
		if err != nil {
			return nil, err
		}
	}
	return doc, nil
}

func applySet(doc, operator bson.M) (bson.M, error) {
	for k, v := range operator {
		if k == entity.Name {
			return nil, errModifiedReserved
		}
		doc[k] = v
	}
	return doc, nil
}