323 lines
6.8 KiB
Go
323 lines
6.8 KiB
Go
package driver
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"local/dndex/config"
|
|
"local/dndex/storage/entity"
|
|
"os"
|
|
"regexp"
|
|
"strings"
|
|
"time"
|
|
|
|
"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, &bolt.Options{
|
|
Timeout: time.Second * 3,
|
|
})
|
|
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
|
|
case "$unset":
|
|
op = applyUnset
|
|
default:
|
|
return nil, errors.New("cannot apply operation " + k)
|
|
}
|
|
doc, err = op(doc, operateBM)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
return doc, nil
|
|
}
|
|
|
|
func applyUnset(doc, operator bson.M) (bson.M, error) {
|
|
for k := range operator {
|
|
if k == entity.Name {
|
|
return nil, errModifiedReserved
|
|
}
|
|
nesting := strings.Split(k, ".")
|
|
if len(nesting) > 1 {
|
|
mInterface, ok := doc[nesting[0]]
|
|
if !ok {
|
|
return nil, fmt.Errorf("path does not exist: %s (%s): %+v", k, nesting[0], doc)
|
|
}
|
|
m, ok := mInterface.(map[string]interface{})
|
|
if !ok {
|
|
pm, pmok := mInterface.(primitive.M)
|
|
m = map[string]interface{}(pm)
|
|
ok = pmok
|
|
}
|
|
if !ok {
|
|
return nil, fmt.Errorf("subpath cannot be followed for non object: %s (%s): %+v (%T)", k, nesting[0], mInterface, mInterface)
|
|
}
|
|
subdoc, err := applyUnset(bson.M(m), bson.M{strings.Join(nesting[1:], "."): ""})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
operator[k] = subdoc
|
|
} else {
|
|
delete(doc, k)
|
|
}
|
|
}
|
|
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
|
|
}
|