From 0124c364730e83ab6c025eed15cdf7dd39e95e4d Mon Sep 17 00:00:00 2001
From: Stas Dmytryshyn <iqpirat@gmail.com>
Date: Mon, 16 Oct 2023 23:58:38 +0200
Subject: [PATCH] feat: expose mongo read pref (#1485)

---
 pkg/storage/mongodb/client.go      | 12 ++++++++++--
 pkg/storage/mongodb/client_test.go |  2 ++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/pkg/storage/mongodb/client.go b/pkg/storage/mongodb/client.go
index 23eefaf83..8049350dc 100644
--- a/pkg/storage/mongodb/client.go
+++ b/pkg/storage/mongodb/client.go
@@ -31,7 +31,8 @@ type Client struct {
 
 func New(connString string, databaseName string, opts ...ClientOpt) (*Client, error) {
 	op := &clientOpts{
-		timeout: defaultTimeout,
+		timeout:  defaultTimeout,
+		readPref: readpref.Nearest(),
 	}
 
 	for _, fn := range opts {
@@ -41,7 +42,7 @@ func New(connString string, databaseName string, opts ...ClientOpt) (*Client, er
 	mongoOpts := mongooptions.Client()
 	mongoOpts.ApplyURI(connString)
 	mongoOpts.SetWriteConcern(writeconcern.New(writeconcern.WMajority(), writeconcern.WTimeout(op.timeout)))
-	mongoOpts.ReadPreference = readpref.Nearest()
+	mongoOpts.ReadPreference = op.readPref
 
 	if op.traceProvider != nil {
 		mongoOpts.Monitor = otelmongo.NewMonitor(otelmongo.WithTracerProvider(op.traceProvider))
@@ -94,6 +95,7 @@ func (c *Client) Close() error {
 type clientOpts struct {
 	timeout       time.Duration
 	traceProvider trace.TracerProvider
+	readPref      *readpref.ReadPref
 }
 
 type ClientOpt func(opts *clientOpts)
@@ -104,6 +106,12 @@ func WithTimeout(timeout time.Duration) ClientOpt {
 	}
 }
 
+func WithReadPref(pref *readpref.ReadPref) ClientOpt {
+	return func(opts *clientOpts) {
+		opts.readPref = pref
+	}
+}
+
 func WithTraceProvider(traceProvider trace.TracerProvider) ClientOpt {
 	return func(opts *clientOpts) {
 		opts.traceProvider = traceProvider
diff --git a/pkg/storage/mongodb/client_test.go b/pkg/storage/mongodb/client_test.go
index ed5529859..2b1c4a629 100644
--- a/pkg/storage/mongodb/client_test.go
+++ b/pkg/storage/mongodb/client_test.go
@@ -21,6 +21,7 @@ import (
 	"go.mongodb.org/mongo-driver/bson/bsontype"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/mongo/readpref"
 	"go.opentelemetry.io/otel/trace"
 
 	"github.com/trustbloc/vcs/pkg/storage/mongodb"
@@ -42,6 +43,7 @@ func TestClient(t *testing.T) {
 
 	client, err := mongodb.New(mongoDBConnString, testDatabaseName,
 		mongodb.WithTimeout(testTimeout),
+		mongodb.WithReadPref(readpref.PrimaryPreferred()),
 		mongodb.WithTraceProvider(trace.NewNoopTracerProvider()),
 	)
 	require.NoError(t, err)