diff --git a/main.go b/main.go index 393f080..e7da07f 100644 --- a/main.go +++ b/main.go @@ -16,6 +16,7 @@ type MongoDbConnector interface { GetAll(collectionName string, findPredicate bson.M, structToDeserialize interface{}) error GetOneProject(collectionName string, findPredicate bson.M, projectFields bson.M, structToDeserialize interface{}) error GetAllProject(collectionName string, findPredicate bson.M, projectFields bson.M, structToDeserialize interface{}) error + GetIterator(collectionName string, findPredicate bson.M, opts *mongo.MongoOptions) *mongo.MongoIterator } type CreateMongoConnector MongoDbConnector diff --git a/mongo/mongo.go b/mongo/mongo.go index ba642f6..439f006 100644 --- a/mongo/mongo.go +++ b/mongo/mongo.go @@ -19,6 +19,15 @@ type MongoConnector struct { session *MgoSession } +type MongoIterator struct { + *mgo.Iter + *mgo.Session +} + +type MongoOptions struct { + Projection *bson.M +} + type MgoSession struct { session *mgo.Session isConnected bool @@ -130,3 +139,24 @@ func (c *MongoConnector) Disconnect() { c.session.session.Close() c.session.isConnected = false } + +func (c *MongoConnector) GetIterator(collectionName string, findPredicate bson.M, opts *MongoOptions) *MongoIterator { + sessionCopy := c.session.session.Copy() + collection := sessionCopy.DB(c.session.dataBaseName).C(collectionName) + query := collection.Find(findPredicate) + if opts != nil { + if opts.Projection != nil { + query = query.Select(*opts.Projection) + } + } + return &MongoIterator{query.Iter(), sessionCopy} +} + +func (iter *MongoIterator) Close() error { + err := iter.Iter.Close() + if err != nil { + return err + } + iter.Session.Close() + return nil +}