From 2596a9e630eaec92542c027f96e3a50c4b365402 Mon Sep 17 00:00:00 2001
From: Agniva De Sarker <agnivade@yahoo.co.in>
Date: Tue, 16 Feb 2021 14:00:03 +0530
Subject: [PATCH] MM-32893: Refactor the getting of DB handle (#16943)

* MM-32893: Refactor the getting of DB handle

We add a method in *SqlStore that returns the correct DB handle
depending on the context.

https://mattermost.atlassian.net/browse/MM-32893

```release-note
NONE
```

* fix receiver name
---
 store/sqlstore/context.go    | 10 ++++++++++
 store/sqlstore/team_store.go | 10 +---------
 store/sqlstore/user_store.go | 34 ++++------------------------------
 3 files changed, 15 insertions(+), 39 deletions(-)

diff --git a/store/sqlstore/context.go b/store/sqlstore/context.go
index 1bd41e6a7a..1a6f8ed113 100644
--- a/store/sqlstore/context.go
+++ b/store/sqlstore/context.go
@@ -5,6 +5,8 @@ package sqlstore
 
 import (
 	"context"
+
+	"github.com/mattermost/gorp"
 )
 
 // storeContextKey is the base type for all context keys for the store.
@@ -32,3 +34,11 @@ func hasMaster(ctx context.Context) bool {
 	}
 	return false
 }
+
+// DBFromContext is a helper utility that returns the DB handle from a given context.
+func (ss *SqlStore) DBFromContext(ctx context.Context) *gorp.DbMap {
+	if hasMaster(ctx) {
+		return ss.GetMaster()
+	}
+	return ss.GetReplica()
+}
diff --git a/store/sqlstore/team_store.go b/store/sqlstore/team_store.go
index 812939685a..ff91903ecd 100644
--- a/store/sqlstore/team_store.go
+++ b/store/sqlstore/team_store.go
@@ -1157,15 +1157,7 @@ func (s SqlTeamStore) GetTeamsForUser(ctx context.Context, userId string) ([]*mo
 	}
 
 	var dbMembers teamMemberWithSchemeRolesList
-
-	var db *gorp.DbMap
-	if hasMaster(ctx) {
-		db = s.GetMaster()
-	} else {
-		db = s.GetReplica()
-	}
-
-	_, err = db.Select(&dbMembers, queryString, args...)
+	_, err = s.SqlStore.DBFromContext(ctx).Select(&dbMembers, queryString, args...)
 	if err != nil {
 		return nil, errors.Wrapf(err, "failed to find TeamMembers with userId=%s", userId)
 	}
diff --git a/store/sqlstore/user_store.go b/store/sqlstore/user_store.go
index 2f93d088aa..f2c669117e 100644
--- a/store/sqlstore/user_store.go
+++ b/store/sqlstore/user_store.go
@@ -335,15 +335,8 @@ func (us SqlUserStore) GetMany(ctx context.Context, ids []string) ([]*model.User
 		return nil, errors.Wrap(err, "users_get_many_tosql")
 	}
 
-	var db *gorp.DbMap
-	if hasMaster(ctx) {
-		db = us.GetMaster()
-	} else {
-		db = us.GetReplica()
-	}
-
 	var users []*model.User
-	if _, err := db.Select(&users, queryString, args...); err != nil {
+	if _, err := us.SqlStore.DBFromContext(ctx).Select(&users, queryString, args...); err != nil {
 		return nil, errors.Wrap(err, "users_get_many_select")
 	}
 
@@ -356,13 +349,7 @@ func (us SqlUserStore) Get(ctx context.Context, id string) (*model.User, error)
 	if err != nil {
 		return nil, errors.Wrap(err, "users_get_tosql")
 	}
-	var db *gorp.DbMap
-	if hasMaster(ctx) {
-		db = us.GetMaster()
-	} else {
-		db = us.GetReplica()
-	}
-	row := db.Db.QueryRow(queryString, args...)
+	row := us.SqlStore.DBFromContext(ctx).Db.QueryRow(queryString, args...)
 
 	var user model.User
 	var props, notifyProps, timezone []byte
@@ -729,15 +716,9 @@ func (us SqlUserStore) GetAllProfilesInChannel(ctx context.Context, channelID st
 	if err != nil {
 		return nil, errors.Wrap(err, "get_all_profiles_in_channel_tosql")
 	}
-	var db *gorp.DbMap
-	if hasMaster(ctx) {
-		db = us.GetMaster()
-	} else {
-		db = us.GetReplica()
-	}
 
 	var users []*model.User
-	rows, err := db.Db.Query(queryString, args...)
+	rows, err := us.SqlStore.DBFromContext(ctx).Db.Query(queryString, args...)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to find Users")
 	}
@@ -961,14 +942,7 @@ func (us SqlUserStore) GetProfileByIds(ctx context.Context, userIds []string, op
 		return nil, errors.Wrap(err, "get_profile_by_ids_tosql")
 	}
 
-	var db *gorp.DbMap
-	if hasMaster(ctx) {
-		db = us.GetMaster()
-	} else {
-		db = us.GetReplica()
-	}
-
-	if _, err := db.Select(&users, queryString, args...); err != nil {
+	if _, err := us.SqlStore.DBFromContext(ctx).Select(&users, queryString, args...); err != nil {
 		return nil, errors.Wrap(err, "failed to find Users")
 	}