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.

228 lines
6.0 KiB

2 months ago
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
}