You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
227 lines
6.0 KiB
227 lines
6.0 KiB
package mongodb
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"go.mongodb.org/mongo-driver/mongo/readpref"
|
|
"net/url"
|
|
"time"
|
|
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
type Storage struct {
|
|
*mongoDatabase
|
|
*mongoCollection
|
|
*mongoClient
|
|
}
|
|
|
|
type Client interface {
|
|
Database(string) Database
|
|
Disconnect(context.Context) error
|
|
StartSession() (mongo.Session, error)
|
|
UseSession(ctx context.Context, fn func(mongo.SessionContext) error) error
|
|
Ping(context.Context) error
|
|
}
|
|
|
|
type Database interface {
|
|
Collection(string) Collection
|
|
}
|
|
|
|
type Collection interface {
|
|
FindOne(context.Context, interface{}) SingleResult
|
|
InsertOne(context.Context, interface{}) (interface{}, error)
|
|
InsertMany(context.Context, []interface{}) ([]interface{}, error)
|
|
DeleteOne(context.Context, interface{}) (int64, error)
|
|
Find(context.Context, interface{}, ...*options.FindOptions) (Cursor, error)
|
|
CountDocuments(context.Context, interface{}, ...*options.CountOptions) (int64, error)
|
|
Aggregate(context.Context, interface{}) (Cursor, error)
|
|
UpdateOne(context.Context, interface{}, interface{}, ...*options.UpdateOptions) (*mongo.UpdateResult, error)
|
|
UpdateMany(context.Context, interface{}, interface{}, ...*options.UpdateOptions) (*mongo.UpdateResult, error)
|
|
}
|
|
|
|
type SingleResult interface {
|
|
Decode(interface{}) error
|
|
}
|
|
|
|
type Cursor interface {
|
|
Close(context.Context) error
|
|
Next(context.Context) bool
|
|
Decode(interface{}) error
|
|
All(context.Context, interface{}) error
|
|
}
|
|
|
|
type mongoClient struct {
|
|
cl *mongo.Client
|
|
}
|
|
|
|
type mongoDatabase struct {
|
|
db *mongo.Database
|
|
}
|
|
|
|
type mongoCollection struct {
|
|
coll *mongo.Collection
|
|
}
|
|
|
|
type mongoSingleResult struct {
|
|
sr *mongo.SingleResult
|
|
}
|
|
|
|
type mongoCursor struct {
|
|
mc *mongo.Cursor
|
|
}
|
|
|
|
type mongoSession struct {
|
|
mongo.Session
|
|
}
|
|
|
|
func (mc *mongoClient) Ping(ctx context.Context) error {
|
|
return mc.cl.Ping(ctx, readpref.Primary())
|
|
}
|
|
|
|
func (mc *mongoClient) Database(dbName string) Database {
|
|
db := mc.cl.Database(dbName)
|
|
return &mongoDatabase{db: db}
|
|
}
|
|
|
|
func (mc *mongoClient) UseSession(ctx context.Context, fn func(mongo.SessionContext) error) error {
|
|
return mc.cl.UseSession(ctx, fn)
|
|
}
|
|
|
|
func (mc *mongoClient) StartSession() (mongo.Session, error) {
|
|
session, err := mc.cl.StartSession()
|
|
return &mongoSession{session}, err
|
|
}
|
|
|
|
func (mc *mongoClient) Disconnect(ctx context.Context) error {
|
|
return mc.cl.Disconnect(ctx)
|
|
}
|
|
|
|
func (md *mongoDatabase) Collection(colName string) Collection {
|
|
collection := md.db.Collection(colName)
|
|
return &mongoCollection{coll: collection}
|
|
}
|
|
|
|
func (mc *mongoCollection) FindOne(ctx context.Context, filter interface{}) SingleResult {
|
|
singleResult := mc.coll.FindOne(ctx, filter)
|
|
return &mongoSingleResult{sr: singleResult}
|
|
}
|
|
|
|
func (mc *mongoCollection) UpdateOne(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
|
|
return mc.coll.UpdateOne(ctx, filter, update, opts[:]...)
|
|
}
|
|
|
|
func (mc *mongoCollection) InsertOne(ctx context.Context, document interface{}) (interface{}, error) {
|
|
id, err := mc.coll.InsertOne(ctx, document)
|
|
return id.InsertedID, err
|
|
}
|
|
|
|
func (mc *mongoCollection) InsertMany(ctx context.Context, document []interface{}) ([]interface{}, error) {
|
|
res, err := mc.coll.InsertMany(ctx, document)
|
|
return res.InsertedIDs, err
|
|
}
|
|
|
|
func (mc *mongoCollection) DeleteOne(ctx context.Context, filter interface{}) (int64, error) {
|
|
count, err := mc.coll.DeleteOne(ctx, filter)
|
|
return count.DeletedCount, err
|
|
}
|
|
|
|
func (mc *mongoCollection) Find(ctx context.Context, filter interface{}, opts ...*options.FindOptions) (Cursor, error) {
|
|
findResult, err := mc.coll.Find(ctx, filter, opts...)
|
|
return &mongoCursor{mc: findResult}, err
|
|
}
|
|
|
|
func (mc *mongoCollection) Aggregate(ctx context.Context, pipeline interface{}) (Cursor, error) {
|
|
aggregateResult, err := mc.coll.Aggregate(ctx, pipeline)
|
|
return &mongoCursor{mc: aggregateResult}, err
|
|
}
|
|
|
|
func (mc *mongoCollection) UpdateMany(ctx context.Context, filter interface{}, update interface{}, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
|
|
return mc.coll.UpdateMany(ctx, filter, update, opts[:]...)
|
|
}
|
|
|
|
func (mc *mongoCollection) CountDocuments(ctx context.Context, filter interface{}, opts ...*options.CountOptions) (int64, error) {
|
|
return mc.coll.CountDocuments(ctx, filter, opts...)
|
|
}
|
|
|
|
func (sr *mongoSingleResult) Decode(v interface{}) error {
|
|
return sr.sr.Decode(v)
|
|
}
|
|
|
|
func (mr *mongoCursor) Close(ctx context.Context) error {
|
|
return mr.mc.Close(ctx)
|
|
}
|
|
|
|
func (mr *mongoCursor) Next(ctx context.Context) bool {
|
|
return mr.mc.Next(ctx)
|
|
}
|
|
|
|
func (mr *mongoCursor) Decode(v interface{}) error {
|
|
return mr.mc.Decode(v)
|
|
}
|
|
|
|
func (mr *mongoCursor) All(ctx context.Context, result interface{}) error {
|
|
return mr.mc.All(ctx, result)
|
|
}
|
|
|
|
// New creates a new MongoDB storage
|
|
func New(config ...Config) *Storage {
|
|
// Set default config
|
|
cfg := configDefault(config...)
|
|
|
|
// Create data source name
|
|
var dsn string
|
|
|
|
// Check if user supplied connection string
|
|
if cfg.ConnectionURI != "" {
|
|
dsn = cfg.ConnectionURI
|
|
} else {
|
|
dsn = "mongodb://"
|
|
if cfg.Username != "" {
|
|
dsn += url.QueryEscape(cfg.Username)
|
|
}
|
|
if cfg.Password != "" {
|
|
dsn += ":" + cfg.Password
|
|
}
|
|
if cfg.Username != "" || cfg.Password != "" {
|
|
dsn += "@"
|
|
}
|
|
dsn += fmt.Sprintf("%s:%d", url.QueryEscape(cfg.Host), cfg.Port)
|
|
}
|
|
|
|
// Set mongo options
|
|
opt := options.Client().ApplyURI(dsn)
|
|
|
|
// Create and connect the mongo client in one step
|
|
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
|
|
defer cancel()
|
|
|
|
client, err := mongo.Connect(ctx, opt)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// verify that the client can connect
|
|
if err = client.Ping(context.Background(), nil); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// Get collection from database
|
|
db := client.Database(cfg.Database)
|
|
col := db.Collection(cfg.Collection)
|
|
|
|
if cfg.Reset {
|
|
if err = col.Drop(context.Background()); err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
store := &Storage{
|
|
mongoDatabase: &mongoDatabase{db: db},
|
|
mongoCollection: &mongoCollection{coll: col},
|
|
mongoClient: &mongoClient{cl: client},
|
|
}
|
|
|
|
return store
|
|
}
|
|
|