diff --git a/mongo/mongo.go b/mongo/mongo.go index 058c407..68ac502 100644 --- a/mongo/mongo.go +++ b/mongo/mongo.go @@ -154,7 +154,8 @@ func (m *Mongo) RoundTrip(msg *Message, tags []string) (_ *Message, err error) { requestCursorID, _ := msg.Op.CursorID() requestCommand, collection := msg.Op.CommandAndCollection() transactionDetails := msg.Op.TransactionDetails() - server, err := m.selectServer(requestCursorID, collection, transactionDetails) + readPref, _ := msg.Op.ReadPref() + server, err := m.selectServer(requestCursorID, collection, transactionDetails, readPref) if err != nil { return nil, err } @@ -230,7 +231,7 @@ func (m *Mongo) RoundTrip(msg *Message, tags []string) (_ *Message, err error) { }, nil } -func (m *Mongo) selectServer(requestCursorID int64, collection string, transDetails *TransactionDetails) (server driver.Server, err error) { +func (m *Mongo) selectServer(requestCursorID int64, collection string, transDetails *TransactionDetails, rp *readpref.ReadPref) (server driver.Server, err error) { defer func(start time.Time) { _ = m.statsd.Timing("server_selection", time.Since(start), []string{fmt.Sprintf("success:%v", err == nil)}, 1) }(time.Now()) @@ -252,9 +253,13 @@ func (m *Mongo) selectServer(requestCursorID int64, collection string, transDeta } } - // Select a server + // Select a server using the provided read preference + // If no read preference was provided, default to primary + if rp == nil { + rp = readpref.Primary() + } selector := description.CompositeSelector([]description.ServerSelector{ - description.ReadPrefSelector(readpref.Primary()), // ignored by sharded clusters + description.ReadPrefSelector(rp), // use client's read preference (ignored by sharded clusters) description.LatencySelector(15 * time.Millisecond), // default localThreshold for the client }) return m.topology.SelectServer(m.roundTripCtx, selector) diff --git a/mongo/operations.go b/mongo/operations.go index d105c9c..9cc120b 100644 --- a/mongo/operations.go +++ b/mongo/operations.go @@ -6,6 +6,7 @@ import ( "fmt" "strings" + "go.mongodb.org/mongo-driver/mongo/readpref" "go.mongodb.org/mongo-driver/x/bsonx/bsoncore" "go.mongodb.org/mongo-driver/x/mongo/driver" "go.mongodb.org/mongo-driver/x/mongo/driver/wiremessage" @@ -33,6 +34,7 @@ type Operation interface { Unacknowledged() bool CommandAndCollection() (Command, string) TransactionDetails() *TransactionDetails + ReadPref() (rp *readpref.ReadPref, ok bool) } // see https://github.com/mongodb/mongo-go-driver/blob/v1.7.2/x/mongo/driver/operation.go#L1361-L1426 @@ -121,6 +123,10 @@ func (o *opUnknown) String() string { return fmt.Sprintf("{ OpUnknown opCode: %d, wm: %s }", o.opCode, o.wm) } +func (o *opUnknown) ReadPref() (rp *readpref.ReadPref, ok bool) { + return readpref.Primary(), false +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#wire-op-query type opQuery struct { reqID int32 @@ -230,6 +236,10 @@ func (q *opQuery) String() string { return fmt.Sprintf("{ OpQuery flags: %s, fullCollectionName: %s, numberToSkip: %d, numberToReturn: %d, query: %s, returnFieldsSelector: %s }", q.flags.String(), q.fullCollectionName, q.numberToSkip, q.numberToReturn, q.query.String(), q.returnFieldsSelector.String()) } +func (q *opQuery) ReadPref() (rp *readpref.ReadPref, ok bool) { + return extractReadPref(q.query) +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op-msg type opMsg struct { reqID int32 @@ -485,6 +495,17 @@ func (m *opMsg) String() string { return fmt.Sprintf("{ OpMsg flags: %d, sections: [%s], checksum: %d }", m.flags, strings.Join(sections, ", "), m.checksum) } +func (m *opMsg) ReadPref() (rp *readpref.ReadPref, ok bool) { + if len(m.sections) == 0 { + return readpref.Primary(), false + } + single, ok := m.sections[0].(*opMsgSectionSingle) + if !ok { + return readpref.Primary(), false + } + return extractReadPref(single.msg) +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op-reply type opReply struct { reqID int32 @@ -588,6 +609,11 @@ func (r *opReply) String() string { return fmt.Sprintf("{ OpReply flags: %d, cursorID: %d, startingFrom: %d, numReturned: %d, documents: [%s] }", r.flags, r.cursorID, r.startingFrom, r.numReturned, strings.Join(documents, ", ")) } +func (r *opReply) ReadPref() (rp *readpref.ReadPref, ok bool) { + // Replies don't contain read preferences + return readpref.Primary(), false +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op-get-more type opGetMore struct { reqID int32 @@ -676,6 +702,11 @@ func (g *opGetMore) String() string { return fmt.Sprintf("{ OpGetMore fullCollectionName: %s, numberToReturn: %d, cursorID: %d }", g.fullCollectionName, g.numberToReturn, g.cursorID) } +func (g *opGetMore) ReadPref() (rp *readpref.ReadPref, ok bool) { + // GetMore operations don't specify read preference, they inherit from cursor + return readpref.Primary(), false +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_update type opUpdate struct { reqID int32 @@ -761,6 +792,11 @@ func (u *opUpdate) String() string { return fmt.Sprintf("{ OpQuery fullCollectionName: %s, flags: %d, selector: %s, update: %s }", u.fullCollectionName, u.flags, u.selector.String(), u.update.String()) } +func (u *opUpdate) ReadPref() (rp *readpref.ReadPref, ok bool) { + // Update is a write operation, always goes to primary + return readpref.Primary(), false +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_insert type opInsert struct { reqID int32 @@ -845,6 +881,11 @@ func (i *opInsert) String() string { return fmt.Sprintf("{ OpInsert flags: %d, fullCollectionName: %s, documents: %s }", i.flags, i.fullCollectionName, strings.Join(documents, ", ")) } +func (i *opInsert) ReadPref() (rp *readpref.ReadPref, ok bool) { + // Insert is a write operation, always goes to primary + return readpref.Primary(), false +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_insert type opDelete struct { reqID int32 @@ -928,6 +969,11 @@ func (d *opDelete) String() string { return fmt.Sprintf("{ OpDelete fullCollectionName: %s, flags: %d, selector: %s }", d.fullCollectionName, d.flags, d.selector.String()) } +func (d *opDelete) ReadPref() (rp *readpref.ReadPref, ok bool) { + // Delete is a write operation, always goes to primary + return readpref.Primary(), false +} + // https://docs.mongodb.com/manual/reference/mongodb-wire-protocol/#op_kill_cursors type opKillCursors struct { reqID int32 @@ -1006,6 +1052,11 @@ func (k *opKillCursors) String() string { return fmt.Sprintf("{ OpKillCursors cursorIDs: %v }", k.cursorIDs) } +func (k *opKillCursors) ReadPref() (rp *readpref.ReadPref, ok bool) { + // KillCursors doesn't use read preferences + return readpref.Primary(), false +} + func appendi32(dst []byte, i32 int32) []byte { return append(dst, byte(i32), byte(i32>>8), byte(i32>>16), byte(i32>>24)) } @@ -1030,3 +1081,49 @@ func readCString(src []byte) (string, []byte, bool) { } return string(src[:idx]), src[idx+1:], true } + +// extractReadPref extracts the read preference from a BSON document +// MongoDB clients send read preferences in the $readPreference field +// See: https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.rst#passing-read-preference-to-mongos +func extractReadPref(doc bsoncore.Document) (*readpref.ReadPref, bool) { + // Check for $readPreference field + rpVal := doc.Lookup("$readPreference") + if rpVal.Type == 0 { + // No read preference specified, default to primary + return readpref.Primary(), false + } + + rpDoc, ok := rpVal.DocumentOK() + if !ok { + return readpref.Primary(), false + } + + // Extract mode + modeVal := rpDoc.Lookup("mode") + mode, ok := modeVal.StringValueOK() + if !ok { + return readpref.Primary(), false + } + + // Parse the mode string and create appropriate read preference + var rp *readpref.ReadPref + switch mode { + case "primary": + rp = readpref.Primary() + case "primaryPreferred": + rp = readpref.PrimaryPreferred() + case "secondary": + rp = readpref.Secondary() + case "secondaryPreferred": + rp = readpref.SecondaryPreferred() + case "nearest": + rp = readpref.Nearest() + default: + return readpref.Primary(), false + } + + // TODO: Support tag sets and maxStalenessSeconds if needed + // For now, we just use the mode without additional options + + return rp, true +} diff --git a/mongo/readpref_test.go b/mongo/readpref_test.go new file mode 100644 index 0000000..4df218f --- /dev/null +++ b/mongo/readpref_test.go @@ -0,0 +1,224 @@ +package mongo + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "go.mongodb.org/mongo-driver/bson" + "go.mongodb.org/mongo-driver/mongo/readpref" + "go.mongodb.org/mongo-driver/x/bsonx/bsoncore" +) + +func TestExtractReadPref_Primary(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "primary"}}}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.True(t, ok, "Should successfully extract read preference") + assert.NotNil(t, rp) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) +} + +func TestExtractReadPref_Secondary(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "secondary"}}}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.True(t, ok, "Should successfully extract read preference") + assert.NotNil(t, rp) + assert.Equal(t, readpref.SecondaryMode, rp.Mode()) +} + +func TestExtractReadPref_SecondaryPreferred(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "secondaryPreferred"}}}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.True(t, ok, "Should successfully extract read preference") + assert.NotNil(t, rp) + assert.Equal(t, readpref.SecondaryPreferredMode, rp.Mode()) +} + +func TestExtractReadPref_PrimaryPreferred(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "primaryPreferred"}}}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.True(t, ok, "Should successfully extract read preference") + assert.NotNil(t, rp) + assert.Equal(t, readpref.PrimaryPreferredMode, rp.Mode()) +} + +func TestExtractReadPref_Nearest(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "nearest"}}}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.True(t, ok, "Should successfully extract read preference") + assert.NotNil(t, rp) + assert.Equal(t, readpref.NearestMode, rp.Mode()) +} + +func TestExtractReadPref_NoReadPreference(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.False(t, ok, "Should return false when no read preference is specified") + assert.NotNil(t, rp) + // Should default to primary + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) +} + +func TestExtractReadPref_InvalidMode(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "invalidMode"}}}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.False(t, ok, "Should return false for invalid read preference mode") + assert.NotNil(t, rp) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) +} + +func TestExtractReadPref_MalformedDocument(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "test"}, + {Key: "$db", Value: "testdb"}, + {Key: "$readPreference", Value: "notADocument"}, + }) + assert.NoError(t, err) + + rp, ok := extractReadPref(doc) + assert.False(t, ok, "Should return false for malformed read preference") + assert.NotNil(t, rp) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) +} + +func TestOpMsg_ReadPref_Primary(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "trainers"}, + {Key: "$db", Value: "test"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "primary"}}}, + }) + assert.NoError(t, err) + + op := NewOpMsg(doc, []bsoncore.Document{}) + + rp, ok := op.Op.ReadPref() + assert.True(t, ok) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) +} + +func TestOpMsg_ReadPref_Secondary(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "trainers"}, + {Key: "$db", Value: "test"}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "secondary"}}}, + }) + assert.NoError(t, err) + + op := NewOpMsg(doc, []bsoncore.Document{}) + + rp, ok := op.Op.ReadPref() + assert.True(t, ok) + assert.Equal(t, readpref.SecondaryMode, rp.Mode()) +} + +func TestOpMsg_ReadPref_NoPreference(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "find", Value: "trainers"}, + {Key: "$db", Value: "test"}, + }) + assert.NoError(t, err) + + op := NewOpMsg(doc, []bsoncore.Document{}) + + rp, ok := op.Op.ReadPref() + assert.False(t, ok) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) +} + +func TestOpQuery_ReadPref(t *testing.T) { + doc, err := bson.Marshal(bson.D{ + {Key: "$query", Value: bson.D{{Key: "name", Value: "test"}}}, + {Key: "$readPreference", Value: bson.D{{Key: "mode", Value: "secondary"}}}, + }) + assert.NoError(t, err) + + op := &opQuery{ + fullCollectionName: "test.trainers", + query: doc, + } + + rp, ok := op.ReadPref() + assert.True(t, ok) + assert.Equal(t, readpref.SecondaryMode, rp.Mode()) +} + +func TestWriteOperations_AlwaysReturnPrimary(t *testing.T) { + // Test Update + updateOp := &opUpdate{} + rp, ok := updateOp.ReadPref() + assert.False(t, ok) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) + + // Test Insert + insertOp := &opInsert{} + rp, ok = insertOp.ReadPref() + assert.False(t, ok) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) + + // Test Delete + deleteOp := &opDelete{} + rp, ok = deleteOp.ReadPref() + assert.False(t, ok) + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) +} + +// TestGetMoreOperation_CursorPinning verifies that getMore operations +// don't specify read preferences directly, as they inherit the server +// from the cursor cache. This is the correct behavior since cursors +// must remain pinned to the server where they were created. +func TestGetMoreOperation_CursorPinning(t *testing.T) { + getMoreOp := &opGetMore{ + cursorID: 12345, + fullCollectionName: "test.collection", + } + + // GetMore should not have a read preference since it uses cursor pinning + rp, ok := getMoreOp.ReadPref() + assert.False(t, ok, "GetMore should not specify read preference") + assert.Equal(t, readpref.PrimaryMode, rp.Mode()) + + // Verify cursor ID is accessible + cursorID, cursorOK := getMoreOp.CursorID() + assert.True(t, cursorOK) + assert.Equal(t, int64(12345), cursorID) +}