[MM-44329] Add a method to store.Store interface to get sql.DB object (#20351)

* [MM-44329] Added methods to sqlstore.Store

* [MM-44329] Added methods to the store.Store interface

* [MM-44329] Regenerated mocks

* [MM-44329] Added comments for the store.Store interface

Co-authored-by: Tim Scheuermann <tim.scheuermann@mattermost.com>

store/sqlstore/store.go

@@ -381,11 +381,12 @@ func (ss *SqlStore) SetMasterX(db *sql.DB) {
 	}
 }
 
+func (ss *SqlStore) GetInternalMasterDB() *sql.DB {
+	return ss.GetMasterX().DB.DB
+}
+
 func (ss *SqlStore) GetSearchReplicaX() *sqlxDBWrapper {
-	ss.licenseMutex.RLock()
-	license := ss.license
-	ss.licenseMutex.RUnlock()
-	if license == nil {
+	if !ss.hasLicense() {
 		return ss.GetMasterX()
 	}
 
@@ -398,10 +399,7 @@ func (ss *SqlStore) GetSearchReplicaX() *sqlxDBWrapper {
 }
 
 func (ss *SqlStore) GetReplicaX() *sqlxDBWrapper {
-	ss.licenseMutex.RLock()
-	license := ss.license
-	ss.licenseMutex.RUnlock()
-	if len(ss.settings.DataSourceReplicas) == 0 || ss.lockedToMaster || license == nil {
+	if len(ss.settings.DataSourceReplicas) == 0 || ss.lockedToMaster || !ss.hasLicense() {
 		return ss.GetMasterX()
 	}
 
@@ -409,6 +407,21 @@ func (ss *SqlStore) GetReplicaX() *sqlxDBWrapper {
 	return ss.ReplicaXs[rrNum]
 }
 
+func (ss *SqlStore) GetInternalReplicaDBs() []*sql.DB {
+	if len(ss.settings.DataSourceReplicas) == 0 || ss.lockedToMaster || !ss.hasLicense() {
+		return []*sql.DB{
+			ss.GetMasterX().DB.DB,
+		}
+	}
+
+	dbs := make([]*sql.DB, len(ss.ReplicaXs))
+	for i, rx := range ss.ReplicaXs {
+		dbs[i] = rx.DB.DB
+	}
+
+	return dbs
+}
+
 func (ss *SqlStore) TotalMasterDbConnections() int {
 	return ss.GetMasterX().Stats().OpenConnections
 }
@@ -946,6 +959,14 @@ func (ss *SqlStore) GetLicense() *model.License {
 	return ss.license
 }
 
+func (ss *SqlStore) hasLicense() bool {
+	ss.licenseMutex.Lock()
+	hasLicense := ss.license != nil
+	ss.licenseMutex.Unlock()
+
+	return hasLicense
+}
+
 func (ss *SqlStore) migrate(direction migrationDirection) error {
 	assets := db.Assets()
 

store/store.go

@@ -7,6 +7,7 @@ package store
 
 import (
 	"context"
+	"database/sql"
 	"time"
 
 	"github.com/mattermost/mattermost-server/v6/model"
@@ -66,6 +67,12 @@ type Store interface {
 	GetDBSchemaVersion() (int, error)
 	GetAppliedMigrations() ([]model.AppliedMigration, error)
 	GetDbVersion(numerical bool) (string, error)
+	// GetInternalMasterDB allows access to the raw master DB
+	// handle for the multi-product architecture.
+	GetInternalMasterDB() *sql.DB
+	// GetInternalReplicaDBs allows access to the raw replica DB
+	// handles for the multi-product architecture.
+	GetInternalReplicaDBs() []*sql.DB
 	TotalMasterDbConnections() int
 	TotalReadDbConnections() int
 	TotalSearchDbConnections() int

store/storetest/mocks/Store.go

@@ -10,6 +10,8 @@ import (
 	model "github.com/mattermost/mattermost-server/v6/model"
 	mock "github.com/stretchr/testify/mock"
 
+	sql "database/sql"
+
 	store "github.com/mattermost/mattermost-server/v6/store"
 
 	time "time"
@@ -287,6 +289,38 @@ func (_m *Store) GetDbVersion(numerical bool) (string, error) {
 	return r0, r1
 }
 
+// GetInternalMasterDB provides a mock function with given fields:
+func (_m *Store) GetInternalMasterDB() *sql.DB {
+	ret := _m.Called()
+
+	var r0 *sql.DB
+	if rf, ok := ret.Get(0).(func() *sql.DB); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*sql.DB)
+		}
+	}
+
+	return r0
+}
+
+// GetInternalReplicaDBs provides a mock function with given fields:
+func (_m *Store) GetInternalReplicaDBs() []*sql.DB {
+	ret := _m.Called()
+
+	var r0 []*sql.DB
+	if rf, ok := ret.Get(0).(func() []*sql.DB); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]*sql.DB)
+		}
+	}
+
+	return r0
+}
+
 // Group provides a mock function with given fields:
 func (_m *Store) Group() store.GroupStore {
 	ret := _m.Called()

store/storetest/store.go

@@ -5,6 +5,7 @@ package storetest
 
 import (
 	"context"
+	"database/sql"
 	"time"
 
 	"github.com/stretchr/testify/mock"
@@ -103,6 +104,8 @@ func (s *Store) LockToMaster()                           { /* do nothing */ }
 func (s *Store) UnlockFromMaster()                       { /* do nothing */ }
 func (s *Store) DropAllTables()                          { /* do nothing */ }
 func (s *Store) GetDbVersion(bool) (string, error)       { return "", nil }
+func (s *Store) GetInternalMasterDB() *sql.DB            { return nil }
+func (s *Store) GetInternalReplicaDBs() []*sql.DB        { return nil }
 func (s *Store) RecycleDBConnections(time.Duration)      {}
 func (s *Store) GetDBSchemaVersion() (int, error)        { return 1, nil }
 func (s *Store) GetAppliedMigrations() ([]model.AppliedMigration, error) {