package driver import ( "context" "errors" "fmt" "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" ) type BoltDB struct { db *bolt.DB } func NewBoltDB(path string) *BoltDB { db, err := bolt.Open(path, 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.ID]; !ok { return errors.New("primary key required to insert: did not find " + entity.ID) } else if _, ok := m[entity.ID].(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.ID].(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) != "nil" && fmt.Sprint(v) != "" { 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 } 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.ID { continue } nesting := strings.Split(k, ".") if len(nesting) > 1 { mInterface, ok := doc[nesting[0]] if !ok { continue //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 of %v (%v) cannot be followed for non object: %s (%s): %+v (%T)", doc, doc[nesting[0]], 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 if len(k) > 0 { delete(doc, k) } } return doc, nil } func applySet(doc, operator bson.M) (bson.M, error) { for k, v := range operator { if k == entity.ID { continue } nesting := strings.Split(k, ".") if len(nesting) > 1 { mInterface, ok := doc[nesting[0]] if !ok { mInterface = map[string]interface{}{} //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 of %v (%v) cannot be followed for non object: %s (%s): %+v (%T)", doc, doc[nesting[0]], k, nesting[0], mInterface, mInterface) } subdoc, err := applySet(bson.M(m), bson.M{strings.Join(nesting[1:], "."): v}) if err != nil { return nil, err } doc[nesting[0]] = subdoc } else if len(k) > 0 { doc[k] = v } } return doc, nil }