diff --git a/store/sqlstore/store.go b/store/sqlstore/store.go index 7080d12e0c..903df4194c 100644 --- a/store/sqlstore/store.go +++ b/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() diff --git a/store/store.go b/store/store.go index 752f6c9030..e858153d62 100644 --- a/store/store.go +++ b/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 diff --git a/store/storetest/mocks/Store.go b/store/storetest/mocks/Store.go index a8665de19b..05998c1458 100644 --- a/store/storetest/mocks/Store.go +++ b/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() diff --git a/store/storetest/store.go b/store/storetest/store.go index f814f0a9b6..b58625e4d4 100644 --- a/store/storetest/store.go +++ b/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) {