From f955451fe532a4586354656453300f6452c85c39 Mon Sep 17 00:00:00 2001
From: Anton Kucherov <anton.kucherov@percona.com>
Date: Mon, 1 Jul 2019 12:53:15 +0300
Subject: [PATCH] PMM-4278 Replace bson to official mongo-driver.

---
 Gopkg.lock                                    | 25 +++++++++++++
 Gopkg.toml                                    |  4 +++
 src/go/mongolib/explain/explain.go            |  9 ++---
 src/go/mongolib/explain/explain_test.go       |  9 ++---
 .../mongolib/fingerprinter/fingerprinter.go   |  5 +--
 src/go/mongolib/proto/bson.go                 | 25 ++++++-------
 src/go/mongolib/proto/main_test.go            |  2 +-
 src/go/mongolib/proto/oplog.go                |  2 +-
 src/go/mongolib/proto/replconfig.go           |  4 +--
 src/go/mongolib/proto/system.profile.go       | 36 +++++++++----------
 src/go/mongolib/stats/stats.go                |  4 +--
 src/go/mongolib/util/util.go                  | 14 +++++---
 12 files changed, 88 insertions(+), 51 deletions(-)

diff --git a/Gopkg.lock b/Gopkg.lock
index 84e073d9..b8012e32 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -63,6 +63,14 @@
   revision = "97b6244175ae18ea6eef668034fd6565847501c9"
   version = "v1.2.4"
 
+[[projects]]
+  digest = "1:a01080d20c45c031c13f3828c56e58f4f51d926a482ad10cc0316225097eb7ea"
+  name = "github.com/go-stack/stack"
+  packages = ["."]
+  pruneopts = ""
+  revision = "2fee6af1a9795aafbe0253a0cfbdf668e1fb8a9a"
+  version = "v1.8.0"
+
 [[projects]]
   digest = "1:530233672f656641b365f8efb38ed9fba80e420baff2ce87633813ab3755ed6d"
   name = "github.com/golang/mock"
@@ -204,6 +212,21 @@
   revision = "dae0fa8d5b0c810a8ab733fbd5510c7cae84eca4"
   version = "v1.4.0"
 
+[[projects]]
+  digest = "1:52376909e3b93cfc66ff439bdf9a7bf88783a0017ceded23c3f1c76b9f7e10ee"
+  name = "go.mongodb.org/mongo-driver"
+  packages = [
+    "bson",
+    "bson/bsoncodec",
+    "bson/bsonrw",
+    "bson/bsontype",
+    "bson/primitive",
+    "x/bsonx/bsoncore",
+  ]
+  pruneopts = ""
+  revision = "582ff343271e8893d785ff094855498c285bce0a"
+  version = "v1.0.3"
+
 [[projects]]
   branch = "master"
   digest = "1:36ef1d8645934b1744cc7d8726e00d3dd9d8d84c18617bf7367a3a6d532f3370"
@@ -276,6 +299,8 @@
     "github.com/satori/go.uuid",
     "github.com/shirou/gopsutil/process",
     "github.com/sirupsen/logrus",
+    "go.mongodb.org/mongo-driver/bson",
+    "go.mongodb.org/mongo-driver/bson/primitive",
     "golang.org/x/crypto/ssh/terminal",
     "gopkg.in/mgo.v2",
     "gopkg.in/mgo.v2/bson",
diff --git a/Gopkg.toml b/Gopkg.toml
index 4b074d04..728cd8d9 100644
--- a/Gopkg.toml
+++ b/Gopkg.toml
@@ -67,3 +67,7 @@
 [[constraint]]
   branch = "v2"
   name = "gopkg.in/mgo.v2"
+
+[[constraint]]
+  name = "go.mongodb.org/mongo-driver"
+  version = "~1.0.0"
diff --git a/src/go/mongolib/explain/explain.go b/src/go/mongolib/explain/explain.go
index abc576b9..2e1c6099 100644
--- a/src/go/mongolib/explain/explain.go
+++ b/src/go/mongolib/explain/explain.go
@@ -3,9 +3,10 @@ package explain
 import (
 	"fmt"
 
-	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
 	"github.com/percona/pmgo"
-	"gopkg.in/mgo.v2/bson"
+	"go.mongodb.org/mongo-driver/bson"
+
+	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
 )
 
 type explain struct {
@@ -22,7 +23,7 @@ func (e *explain) Explain(db string, query []byte) ([]byte, error) {
 	var err error
 	var eq proto.ExampleQuery
 
-	err = bson.UnmarshalJSON(query, &eq)
+	err = bson.UnmarshalExtJSON(query, true, &eq)
 	if err != nil {
 		return nil, fmt.Errorf("explain: unable to decode query %s: %s", string(query), err)
 	}
@@ -37,7 +38,7 @@ func (e *explain) Explain(db string, query []byte) ([]byte, error) {
 		return nil, err
 	}
 
-	resultJson, err := bson.MarshalJSON(result)
+	resultJson, err := bson.MarshalExtJSON(result, true, true)
 	if err != nil {
 		return nil, fmt.Errorf("explain: unable to encode explain result of %s: %s", string(query), err)
 	}
diff --git a/src/go/mongolib/explain/explain_test.go b/src/go/mongolib/explain/explain_test.go
index 8de53865..4af04983 100644
--- a/src/go/mongolib/explain/explain_test.go
+++ b/src/go/mongolib/explain/explain_test.go
@@ -9,10 +9,11 @@ import (
 	"testing"
 
 	"github.com/Masterminds/semver"
+	"github.com/percona/pmgo"
+	"go.mongodb.org/mongo-driver/bson"
+
 	"github.com/percona/percona-toolkit/src/go/lib/tutil"
 	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
-	"github.com/percona/pmgo"
-	"gopkg.in/mgo.v2/bson"
 )
 
 const (
@@ -137,7 +138,7 @@ func TestExplain(t *testing.T) {
 			if err != nil {
 				t.Fatalf("cannot load sample %s: %s", dir+file.Name(), err)
 			}
-			query, err := bson.MarshalJSON(eq)
+			query, err := bson.MarshalExtJSON(eq, true, true)
 			if err != nil {
 				t.Fatalf("cannot marshal json %s: %s", dir+file.Name(), err)
 			}
@@ -150,7 +151,7 @@ func TestExplain(t *testing.T) {
 
 			if err == nil {
 				result := proto.BsonD{}
-				err = bson.UnmarshalJSON(got, &result)
+				err = bson.UnmarshalExtJSON(got, true, &result)
 				if err != nil {
 					t.Fatalf("cannot unmarshal json explain result: %s", err)
 				}
diff --git a/src/go/mongolib/fingerprinter/fingerprinter.go b/src/go/mongolib/fingerprinter/fingerprinter.go
index e25536f6..b42fd8be 100644
--- a/src/go/mongolib/fingerprinter/fingerprinter.go
+++ b/src/go/mongolib/fingerprinter/fingerprinter.go
@@ -7,9 +7,10 @@ import (
 	"sort"
 	"strings"
 
+	"go.mongodb.org/mongo-driver/bson"
+
 	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
 	"github.com/percona/percona-toolkit/src/go/mongolib/util"
-	"gopkg.in/mgo.v2/bson"
 )
 
 var (
@@ -112,7 +113,7 @@ func (f *Fingerprinter) Fingerprint(doc proto.SystemProfile) (Fingerprint, error
 			break
 		}
 		// first key is operation type
-		op = query[0].Name
+		op = query[0].Key
 		collection, _ = query[0].Value.(string)
 		switch op {
 		case "group":
diff --git a/src/go/mongolib/proto/bson.go b/src/go/mongolib/proto/bson.go
index b8664a72..477410ca 100644
--- a/src/go/mongolib/proto/bson.go
+++ b/src/go/mongolib/proto/bson.go
@@ -6,7 +6,8 @@ import (
 	"fmt"
 	"math"
 
-	"gopkg.in/mgo.v2/bson"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
 type BsonD bson.D
@@ -37,8 +38,8 @@ func (d *BsonD) UnmarshalJSON(data []byte) error {
 			return fmt.Errorf("expected key to be a string but got %s", t)
 		}
 
-		de := bson.DocElem{}
-		de.Name = key
+		de := primitive.E{}
+		de.Key = key
 
 		if !dec.More() {
 			return fmt.Errorf("missing value for key %s", key)
@@ -51,13 +52,13 @@ func (d *BsonD) UnmarshalJSON(data []byte) error {
 		}
 
 		var v BsonD
-		err = bson.UnmarshalJSON(raw, &v)
+		err = bson.UnmarshalExtJSON(raw, true, &v)
 		if err != nil {
 			var v []BsonD
-			err = bson.UnmarshalJSON(raw, &v)
+			err = bson.UnmarshalExtJSON(raw, true, &v)
 			if err != nil {
 				var v interface{}
-				err = bson.UnmarshalJSON(raw, &v)
+				err = bson.UnmarshalExtJSON(raw, true, &v)
 				if err != nil {
 					return err
 				} else {
@@ -98,7 +99,7 @@ func (d BsonD) MarshalJSON() ([]byte, error) {
 		}
 
 		// marshal key
-		key, err := bson.MarshalJSON(v.Name)
+		key, err := bson.MarshalExtJSON(v.Key, true, true)
 		if err != nil {
 			return nil, err
 		}
@@ -119,7 +120,7 @@ func (d BsonD) MarshalJSON() ([]byte, error) {
 			val = append(val, '"')
 		} else {
 			// marshal value
-			val, err = bson.MarshalJSON(v.Value)
+			val, err = bson.MarshalExtJSON(v.Value, true, true)
 			if err != nil {
 				return nil, err
 			}
@@ -142,13 +143,13 @@ func (d BsonD) Map() (m bson.M) {
 	for _, item := range d {
 		switch v := item.Value.(type) {
 		case BsonD:
-			m[item.Name] = v.Map()
+			m[item.Key] = v.Map()
 		case []BsonD:
 			el := []bson.M{}
 			for i := range v {
 				el = append(el, v[i].Map())
 			}
-			m[item.Name] = el
+			m[item.Key] = el
 		case []interface{}:
 			// mgo/bson doesn't expose UnmarshalBSON interface
 			// so we can't create custom bson.Unmarshal()
@@ -158,9 +159,9 @@ func (d BsonD) Map() (m bson.M) {
 					el = append(el, b.Map())
 				}
 			}
-			m[item.Name] = el
+			m[item.Key] = el
 		default:
-			m[item.Name] = item.Value
+			m[item.Key] = item.Value
 		}
 	}
 	return m
diff --git a/src/go/mongolib/proto/main_test.go b/src/go/mongolib/proto/main_test.go
index 45a00723..bc614127 100644
--- a/src/go/mongolib/proto/main_test.go
+++ b/src/go/mongolib/proto/main_test.go
@@ -6,8 +6,8 @@ import (
 	"os"
 	"testing"
 
+	"go.mongodb.org/mongo-driver/bson"
 	mgo "gopkg.in/mgo.v2"
-	"gopkg.in/mgo.v2/bson"
 	"gopkg.in/mgo.v2/dbtest"
 )
 
diff --git a/src/go/mongolib/proto/oplog.go b/src/go/mongolib/proto/oplog.go
index 371f13cd..9ffd9c13 100644
--- a/src/go/mongolib/proto/oplog.go
+++ b/src/go/mongolib/proto/oplog.go
@@ -3,7 +3,7 @@ package proto
 import (
 	"time"
 
-	"gopkg.in/mgo.v2/bson"
+	"go.mongodb.org/mongo-driver/bson"
 )
 
 type OplogEntry struct {
diff --git a/src/go/mongolib/proto/replconfig.go b/src/go/mongolib/proto/replconfig.go
index 8ce984d2..58ab7058 100644
--- a/src/go/mongolib/proto/replconfig.go
+++ b/src/go/mongolib/proto/replconfig.go
@@ -1,7 +1,7 @@
 package proto
 
 import (
-	"gopkg.in/mgo.v2/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
 type ReplicaSetConfigTags map[string]string
@@ -35,7 +35,7 @@ type ReplicaSetConfigSettings struct {
 	ElectionTimeoutMillis   int64                 `bson:"electionTimeoutMillis,omitempty"`   // The time limit in milliseconds for detecting when a replica set’s primary is unreachable.
 	GetLastErrorDefaults    *GetLastErrorDefaults `bson:"getLastErrorDefaults,omitempty"`    // A document that specifies the write concern for the replica set.
 	GetLastErrorModes       *GetLastErrorModes    `bson:"getLastErrorModes,omitempty"`       // A document used to define an extended write concern through the use of members[n].tags.
-	ReplicaSetId            *bson.ObjectId        `bson:"replicaSetId,omitempty"`            // Replset Id (ObjectId)
+	ReplicaSetId            *primitive.ObjectID   `bson:"replicaSetId,omitempty"`            // Replset Id (ObjectId)
 }
 
 type ReplicaSetConfig struct {
diff --git a/src/go/mongolib/proto/system.profile.go b/src/go/mongolib/proto/system.profile.go
index cfd0e263..dd75e1ac 100644
--- a/src/go/mongolib/proto/system.profile.go
+++ b/src/go/mongolib/proto/system.profile.go
@@ -4,7 +4,7 @@ import (
 	"strings"
 	"time"
 
-	"gopkg.in/mgo.v2/bson"
+	"go.mongodb.org/mongo-driver/bson"
 )
 
 // docsExamined is renamed from nscannedObjects in 3.2.0
@@ -143,9 +143,9 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 			break
 		}
 
-		if cmd.Len() == 0 || cmd[0].Name != "find" {
+		if cmd.Len() == 0 || cmd[0].Key != "find" {
 			var filter interface{}
-			if cmd.Len() > 0 && cmd[0].Name == "query" {
+			if cmd.Len() > 0 && cmd[0].Key == "query" {
 				filter = cmd[0].Value
 			} else {
 				filter = cmd
@@ -163,7 +163,7 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 			}
 		} else {
 			for i := range cmd {
-				switch cmd[i].Name {
+				switch cmd[i].Key {
 				// PMM-1905: Drop "ntoreturn" if it's negative.
 				case "ntoreturn":
 					// If it's non-negative, then we are fine, continue to next param.
@@ -190,13 +190,13 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 		}
 		if cmd.Len() == 0 {
 			cmd = BsonD{
-				{Name: "q", Value: self.Query},
-				{Name: "u", Value: self.UpdateObj},
+				{Key: "q", Value: self.Query},
+				{Key: "u", Value: self.UpdateObj},
 			}
 		}
 		cmd = BsonD{
-			{Name: "update", Value: coll},
-			{Name: "updates", Value: []interface{}{cmd}},
+			{Key: "update", Value: coll},
+			{Key: "updates", Value: []interface{}{cmd}},
 		}
 	case "remove":
 		s := strings.SplitN(self.Ns, ".", 2)
@@ -206,20 +206,20 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 		}
 		if cmd.Len() == 0 {
 			cmd = BsonD{
-				{Name: "q", Value: self.Query},
+				{Key: "q", Value: self.Query},
 				// we can't determine if limit was 1 or 0 so we assume 0
-				{Name: "limit", Value: 0},
+				{Key: "limit", Value: 0},
 			}
 		}
 		cmd = BsonD{
-			{Name: "delete", Value: coll},
-			{Name: "deletes", Value: []interface{}{cmd}},
+			{Key: "delete", Value: coll},
+			{Key: "deletes", Value: []interface{}{cmd}},
 		}
 	case "insert":
 		if cmd.Len() == 0 {
 			cmd = self.Query
 		}
-		if cmd.Len() == 0 || cmd[0].Name != "insert" {
+		if cmd.Len() == 0 || cmd[0].Key != "insert" {
 			coll := ""
 			s := strings.SplitN(self.Ns, ".", 2)
 			if len(s) == 2 {
@@ -235,7 +235,7 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 			cmd = self.OriginatingCommand
 			for i := range cmd {
 				// drop $db param as it is not supported in MongoDB 3.0
-				if cmd[i].Name == "$db" {
+				if cmd[i].Key == "$db" {
 					if len(cmd)-1 == i {
 						cmd = cmd[:i]
 					} else {
@@ -246,11 +246,11 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 			}
 		} else {
 			cmd = BsonD{
-				{Name: "getmore", Value: ""},
+				{Key: "getmore", Value: ""},
 			}
 		}
 	case "command":
-		if cmd.Len() == 0 || cmd[0].Name != "group" {
+		if cmd.Len() == 0 || cmd[0].Key != "group" {
 			break
 		}
 
@@ -268,7 +268,7 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 				//
 				// The $reduce function shouldn't affect explain execution plan (e.g. what indexes are picked)
 				// so we ignore it for now until we find better way to handle this issue
-				if group[i].Name == "$reduce" {
+				if group[i].Key == "$reduce" {
 					group[i].Value = "{}"
 					cmd[0].Value = group
 					break
@@ -279,7 +279,7 @@ func (self ExampleQuery) ExplainCmd() bson.D {
 
 	return bson.D{
 		{
-			Name:  "explain",
+			Key:   "explain",
 			Value: cmd,
 		},
 	}
diff --git a/src/go/mongolib/stats/stats.go b/src/go/mongolib/stats/stats.go
index b3d56061..34cb9f21 100644
--- a/src/go/mongolib/stats/stats.go
+++ b/src/go/mongolib/stats/stats.go
@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"github.com/montanaflynn/stats"
-	"gopkg.in/mgo.v2/bson"
+	"go.mongodb.org/mongo-driver/bson"
 
 	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
 )
@@ -75,7 +75,7 @@ func (s *Stats) Add(doc proto.SystemProfile) error {
 	}
 	if qiac, ok = s.getQueryInfoAndCounters(key); !ok {
 		query := proto.NewExampleQuery(doc)
-		queryBson, err := bson.MarshalJSON(query)
+		queryBson, err := bson.MarshalExtJSON(query, true, true)
 		if err != nil {
 			return err
 		}
diff --git a/src/go/mongolib/util/util.go b/src/go/mongolib/util/util.go
index 524c613d..1d020e6d 100644
--- a/src/go/mongolib/util/util.go
+++ b/src/go/mongolib/util/util.go
@@ -8,14 +8,15 @@ import (
 	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
 	"github.com/percona/pmgo"
 	"github.com/pkg/errors"
-	mgo "gopkg.in/mgo.v2"
-	"gopkg.in/mgo.v2/bson"
+	"go.mongodb.org/mongo-driver/bson"
+	"gopkg.in/mgo.v2" // TODO: Remove this dependency
 )
 
 var (
 	CANNOT_GET_QUERY_ERROR = errors.New("cannot get query field from the profile document (it is not a map)")
 )
 
+// TODO: Refactor to official mongo-driver.
 func GetReplicasetMembers(dialer pmgo.Dialer, di *pmgo.DialInfo) ([]proto.Members, error) {
 	hostnames, err := GetHostnames(dialer, di)
 	if err != nil {
@@ -77,6 +78,7 @@ func GetReplicasetMembers(dialer pmgo.Dialer, di *pmgo.DialInfo) ([]proto.Member
 	return members, nil
 }
 
+// TODO: Refactor to official mongo-driver.
 func GetHostnames(dialer pmgo.Dialer, di *pmgo.DialInfo) ([]string, error) {
 	hostnames := []string{di.Addrs[0]}
 	di.Direct = true
@@ -184,6 +186,7 @@ func buildHostsListFromShardMap(shardsMap proto.ShardsMap) []string {
 
 // This function is like GetHostnames but it uses listShards instead of getShardMap
 // so it won't include config servers in the returned list
+// TODO: Refactor to official mongo-driver.
 func GetShardedHosts(dialer pmgo.Dialer, di *pmgo.DialInfo) ([]string, error) {
 	hostnames := []string{di.Addrs[0]}
 	session, err := dialer.DialWithInfo(di)
@@ -215,6 +218,7 @@ func getTmpDI(di *pmgo.DialInfo, hostname string) *pmgo.DialInfo {
 	return &tmpdi
 }
 
+// TODO: Refactor to official mongo-driver.
 func GetServerStatus(dialer pmgo.Dialer, di *pmgo.DialInfo, hostname string) (proto.ServerStatus, error) {
 	ss := proto.ServerStatus{}
 
@@ -227,8 +231,8 @@ func GetServerStatus(dialer pmgo.Dialer, di *pmgo.DialInfo, hostname string) (pr
 	session.SetMode(mgo.Monotonic, true)
 
 	query := bson.D{
-		{Name: "serverStatus", Value: 1},
-		{Name: "recordStats", Value: 1},
+		{Key: "serverStatus", Value: 1},
+		{Key: "recordStats", Value: 1},
 	}
 	if err := session.DB("admin").Run(query, &ss); err != nil {
 		return ss, errors.Wrap(err, "GetHostInfo.serverStatus")
@@ -299,7 +303,7 @@ func GetQueryField(doc proto.SystemProfile) (bson.M, error) {
 	}
 
 	// {"ns":"test.system.js","op":"query","query":{"find":"system.js"}}
-	if len(query) == 1 && query[0].Name == "find" {
+	if len(query) == 1 && query[0].Key == "find" {
 		return bson.M{}, nil
 	}