Allow partial results flag to be set

Author: kahuang

session.go

@@ -2912,7 +2912,6 @@ func (p *Pipe) SetMaxTime(d time.Duration) *Pipe {
 	return p
 }
 
-
 // Collation allows to specify language-specific rules for string comparison,
 // such as rules for lettercase and accent marks.
 // When specifying collation, the locale field is mandatory; all other collation
@@ -3577,6 +3576,13 @@ func (q *Query) SetMaxTime(d time.Duration) *Query {
 	return q
 }
 
+func (q *Query) AllowPartial() *Query {
+	q.m.Lock()
+	q.op.AllowPartial()
+	q.m.Unlock()
+	return q
+}
+
 // Snapshot will force the performed query to make use of an available
 // index on the _id field to prevent the same document from being returned
 // more than once in a single iteration. This might happen without this
@@ -3753,23 +3759,24 @@ func prepareFindOp(socket *mongoSocket, op *queryOp, limit int32) bool {
 	}
 
 	find := findCmd{
-		Collection:      op.collection[nameDot+1:],
-		Filter:          op.query,
-		Projection:      op.selector,
-		Sort:            op.options.OrderBy,
-		Skip:            op.skip,
-		Limit:           limit,
-		MaxTimeMS:       op.options.MaxTimeMS,
-		MaxScan:         op.options.MaxScan,
-		Hint:            op.options.Hint,
-		Comment:         op.options.Comment,
-		Snapshot:        op.options.Snapshot,
-		Collation:       op.options.Collation,
-		Tailable:        op.flags&flagTailable != 0,
-		AwaitData:       op.flags&flagAwaitData != 0,
-		OplogReplay:     op.flags&flagLogReplay != 0,
-		NoCursorTimeout: op.flags&flagNoCursorTimeout != 0,
-		ReadConcern:     readLevel{level: op.readConcern},
+		Collection:          op.collection[nameDot+1:],
+		Filter:              op.query,
+		Projection:          op.selector,
+		Sort:                op.options.OrderBy,
+		Skip:                op.skip,
+		Limit:               limit,
+		MaxTimeMS:           op.options.MaxTimeMS,
+		MaxScan:             op.options.MaxScan,
+		Hint:                op.options.Hint,
+		Comment:             op.options.Comment,
+		Snapshot:            op.options.Snapshot,
+		Collation:           op.options.Collation,
+		Tailable:            op.flags&flagTailable != 0,
+		AwaitData:           op.flags&flagAwaitData != 0,
+		OplogReplay:         op.flags&flagLogReplay != 0,
+		NoCursorTimeout:     op.flags&flagNoCursorTimeout != 0,
+		ReadConcern:         readLevel{level: op.readConcern},
+		AllowPartialResults: op.flags&flagPartial != 0,
 	}
 
 	if op.limit < 0 {

socket.go

@@ -68,6 +68,8 @@ const (
 	flagLogReplay
 	flagNoCursorTimeout
 	flagAwaitData
+	flagExhaust
+	flagPartial
 )
 
 type queryOp struct {
@@ -98,6 +100,10 @@ type queryWrapper struct {
 	Collation      *Collation  `bson:"$collation,omitempty"`
 }
 
+func (op *queryOp) AllowPartial() {
+	op.flags |= flagPartial
+}
+
 func (op *queryOp) finalQuery(socket *mongoSocket) interface{} {
 	if op.flags&flagSlaveOk != 0 && socket.ServerInfo().Mongos {
 		var modeName string