package storage import ( "context" "errors" "gitea.inhome.blapointe.com/local/storage/resolve" "strings" "time" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/mongo/readconcern" "go.mongodb.org/mongo-driver/mongo/readpref" "go.mongodb.org/mongo-driver/mongo/writeconcern" ) type Mongo struct { db *mongo.Client } func NewMongo(addr string, auth ...string) (*Mongo, error) { addr = strings.TrimPrefix(addr, "mongodb://") var credentials *options.ClientOptions = nil if len(auth) == 2 { credentials = options.Client().SetAuth(options.Credential{ Username: auth[0], Password: auth[1], }) } db, err := mongo.NewClient( options.Client().ApplyURI("mongodb://"+addr), options.Client().SetReadConcern(readconcern.Local()), options.Client().SetReadPreference(readpref.PrimaryPreferred()), options.Client().SetWriteConcern(writeconcern.New( writeconcern.W(1), writeconcern.J(false), )), credentials, ) if err != nil { return nil, err } ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() if err := db.Connect(ctx); err != nil { return nil, err } if err := db.Ping(ctx, nil); err != nil { return nil, err } if _, err := db.ListDatabaseNames(ctx, &bson.M{}); err != nil { return nil, err } return &Mongo{db: db}, nil } func (mg *Mongo) List(ns []string, limits ...string) ([]string, error) { namespace := resolve.Namespace(ns) limits = resolve.Limits(limits) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() collection := mg.db.Database(resolve.DefaultNamespace).Collection(namespace) filter := bson.M{"_id": bson.M{ "$gte": limits[0], "$lte": limits[1], }} projection := bson.M{"_id": 1} cursor, err := collection.Find(ctx, filter, options.Find().SetProjection(projection)) if err != nil { return nil, err } defer cursor.Close(ctx) keys := []string{} for cursor.Next(ctx) { var elem bson.Raw if err := cursor.Decode(&elem); err != nil { return nil, err } if raw, err := elem.LookupErr("_id"); err != nil { return nil, err } else if s, ok := raw.StringValueOK(); !ok { return nil, errors.New("_id is not a string") } else { keys = append(keys, s) } } return keys, nil } func (mg *Mongo) Get(key string, ns ...string) ([]byte, error) { namespace := resolve.Namespace(ns) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() collection := mg.db.Database(resolve.DefaultNamespace).Collection(namespace) filter := bson.M{"_id": key} cursor, err := collection.Find(ctx, filter) if err != nil { return nil, err } defer cursor.Close(ctx) if !cursor.Next(ctx) { return nil, ErrNotFound } var elem bson.Raw if err := cursor.Decode(&elem); err != nil { return nil, err } raw, err := elem.LookupErr("value") if err != nil { return nil, err } _, b, ok := raw.BinaryOK() if !ok { return nil, ErrNotImpl } return b, nil } func (mg *Mongo) Set(key string, value []byte, ns ...string) error { namespace := resolve.Namespace(ns) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() collection := mg.db.Database(resolve.DefaultNamespace).Collection(namespace) filter := bson.M{"_id": key} document := bson.M{"value": value} _, err := collection.ReplaceOne(ctx, filter, document, options.Replace().SetUpsert(true)) return err } func (mg *Mongo) Close() error { return nil }